Add initial datasource module (#368)

Author: cx1111

tests/test_datasource.py

@@ -0,0 +1,106 @@
+import platform
+import pytest
+
+from wfdb.io import (
+    DataSource,
+    DataSourceType,
+    add_data_source,
+    remove_data_source,
+    reset_data_sources,
+)
+
+from wfdb.io.datasource import _data_sources
+
+LOCAL_PATH = (
+    "C:\\Users\\Public\\data"
+    if platform.system() == "Windows"
+    else "/bigdata/smalldata"
+)
+
+
+class TestDataSource:
+    def test_create_valid_local_ds(self):
+        ds = DataSource(
+            name="localds",
+            ds_type=DataSourceType.LOCAL,
+            uri=LOCAL_PATH,
+        )
+        assert ds
+
+    def test_create_invalid_local_ds(self):
+        with pytest.raises(ValueError):
+            DataSource(
+                name="localds",
+                ds_type=DataSourceType.LOCAL,
+                uri="notabsolute",
+            )
+
+    def test_create_valid_http_ds(self):
+        ds = DataSource(
+            name="httpds",
+            ds_type=DataSourceType.HTTP,
+            uri="http://bigdata.com",
+        )
+        assert ds.uri == "http://bigdata.com"
+
+    def test_create_invalid_http_ds(self):
+        with pytest.raises(ValueError):
+            DataSource(
+                name="httpds",
+                ds_type=DataSourceType.HTTP,
+                uri="www.bigdata.com",
+            )
+
+    def test_add_reset_ds(self):
+        ds = DataSource(
+            name="localds",
+            ds_type=DataSourceType.LOCAL,
+            uri=LOCAL_PATH,
+        )
+        add_data_source(ds)
+        assert len(_data_sources) == 2
+        assert _data_sources[ds.name] == ds
+        # We rely on reset_data_sources for test cleanup.
+        reset_data_sources(keep_pn=True)
+        assert len(_data_sources) == 1
+
+    def test_add_multiple_ds(self):
+        ds1 = DataSource(
+            name="localds",
+            ds_type=DataSourceType.LOCAL,
+            uri=LOCAL_PATH,
+        )
+        add_data_source(ds1)
+        ds2 = DataSource(
+            name="anotherlocalds",
+            ds_type=DataSourceType.LOCAL,
+            uri=LOCAL_PATH,
+        )
+        add_data_source(ds2)
+
+        assert len(_data_sources) == 3
+        assert _data_sources[ds1.name] == ds1
+        assert _data_sources[ds2.name] == ds2
+        reset_data_sources(keep_pn=True)
+
+    def test_remove_ds(self):
+        ds = DataSource(
+            name="localds",
+            ds_type=DataSourceType.LOCAL,
+            uri=LOCAL_PATH,
+        )
+        add_data_source(ds)
+        remove_data_source("localds")
+        assert len(_data_sources) == 1
+
+    def test_unique_ds_names(self):
+        ds = DataSource(
+            name="localds",
+            ds_type=DataSourceType.LOCAL,
+            uri=LOCAL_PATH,
+        )
+        add_data_source(ds)
+        # Cannot set multiple data sources with the same name
+        with pytest.raises(ValueError):
+            add_data_source(ds)
+        reset_data_sources(keep_pn=True)

wfdb/__init__.py

@@ -25,6 +25,18 @@
     get_record_list,
     set_db_index_url,
 )
+
+
+from wfdb.io.datasource import (
+    DataSource,
+    DataSourceType,
+    show_data_sources,
+    add_data_source,
+    remove_data_source,
+    reset_data_sources,
+)
+
+
 from wfdb.plot.plot import plot_items, plot_wfdb, plot_all_records
 
 from wfdb.version import __version__

wfdb/io/__init__.py

@@ -27,3 +27,12 @@
     get_record_list,
     set_db_index_url,
 )
+
+from wfdb.io.datasource import (
+    DataSource,
+    DataSourceType,
+    show_data_sources,
+    add_data_source,
+    remove_data_source,
+    reset_data_sources,
+)

wfdb/io/datasource.py

@@ -0,0 +1,93 @@
+from enum import Enum, auto
+from pathlib import Path
+from typing import Dict
+from urllib.parse import urlparse
+
+
+class DataSourceType(Enum):
+    LOCAL = auto()
+    HTTP = auto()
+
+
+class DataSource:
+    def __init__(self, name: str, ds_type: DataSourceType, uri: str):
+        self.name = name
+        self.ds_type = ds_type
+        self.uri = uri
+
+    def __str__(self):
+        return f"{self.name} : {self.ds_type} : {self.uri}"
+
+    @property
+    def uri(self):
+        return self._uri
+
+    @uri.setter
+    def uri(self, value: str):
+        if self.ds_type == DataSourceType.LOCAL:
+            path = Path(value)
+            if not path.is_absolute():
+                raise ValueError(
+                    "uri field for a LOCAL DataSource must be a valid absolute path"
+                )
+        elif self.ds_type is DataSourceType.HTTP:
+            url = urlparse(value)
+            if not url.netloc:
+                raise ValueError(
+                    "uri field for an HTTP DataSource must be a valid URL"
+                )
+        self._uri = value
+
+
+_PHYSIONET_DATA_SOURCE = DataSource(
+    "physionet",
+    DataSourceType.HTTP,
+    "https://physionet.org/files/",
+)
+
+# Dict of configured data sources
+_data_sources: Dict[str, DataSource] = {"physionet": _PHYSIONET_DATA_SOURCE}
+
+
+def show_data_sources():
+    """
+    Displays all configured data sources
+    """
+    print("Data sources:")
+    for _, ds in _data_sources.items():
+        print(ds)
+
+
+def add_data_source(ds: DataSource):
+    """
+    Add a data source to the set of configured data sources
+    """
+    if ds.name in _data_sources:
+        raise ValueError(
+            f"There is already a configured data source with name: {ds.name}"
+        )
+
+    _data_sources[ds.name] = ds
+
+
+def remove_data_source(ds_name: str):
+    """
+    Remove a data source from the set of configured data sources
+    """
+    del _data_sources[ds_name]
+
+
+def reset_data_sources(keep_pn: bool = False):
+    """
+    Reset all configured data sources
+
+    Parameters
+    ----------
+
+    keep_pn : bool
+        If True, keep the default physionet data source.
+
+    """
+    _data_sources.clear()
+    if keep_pn:
+        _data_sources["physionet"] = _PHYSIONET_DATA_SOURCE