Merge pull request #888 from sfu-db/load_db

feat(eda): add get_db_names

Author: jinglinpeng

dataprep/datasets/__init__.py

@@ -1,6 +1,6 @@
 """This module implements load dataset related functions"""
 
 from ._base import load_dataset, _load_dataset_as_dask, load_db
-from ._base import get_dataset_names
+from ._base import get_dataset_names, get_db_names
 
-__all__ = ["load_dataset", "get_dataset_names", "_load_dataset_as_dask", "load_db"]
+__all__ = ["load_dataset", "get_dataset_names", "_load_dataset_as_dask", "load_db", "get_db_names"]

dataprep/datasets/_base.py

@@ -30,6 +30,26 @@ def get_dataset_names() -> List[str]:
     return datasets
 
 
+def get_db_names() -> List[str]:
+    """
+    Get all available database names. It is all csv file names in 'database' folder.
+
+    Returns
+    -------
+    datasets: list
+        A list of all available dataset names.
+
+    """
+    module_path = dirname(__file__)
+    files = os.listdir(f"{module_path}/database")
+    db_files = list(filter(lambda x: x.endswith(".db"), files))
+
+    # remove suffix csv and get dataset names
+    db_names = list(map(lambda f: os.path.splitext(f)[0], db_files))
+
+    return db_names
+
+
 def _get_dataset_path(name: str) -> str:
     """
     Given a dataset name, output the file path.
@@ -80,20 +100,24 @@ def load_dataset(name: str) -> pd.DataFrame:
     return df
 
 
-def load_db(file_name: str) -> Engine:
+def load_db(name: str) -> Engine:
     """
     Load a database file
 
     Parameters
     ----------
-    file_name: str
+    name: str
         Name of the database file
 
     Returns
     -------
     db_url : str
         SQLite url
     """
+    file_name = name.lower()
+    if not file_name.endswith(".db"):
+        file_name += ".db"
+
     db_file_path = str(
         os.path.join(os.path.dirname(os.path.abspath(__file__)), "database", file_name)
     )

dataprep/tests/datasets/test_datasets.py

@@ -1,16 +1,27 @@
 """
 module for testing the functions inside datasets
 """
-from ...datasets import get_dataset_names, load_dataset
+from ...datasets import get_dataset_names, get_db_names, load_dataset, load_db
 
 
 def test_get_dataset_names() -> None:
     names = get_dataset_names()
     assert len(names) > 0
 
 
+def test_get_db_names() -> None:
+    names = get_db_names()
+    assert len(names) > 0
+
+
 def test_load_dataset() -> None:
     dataset_names = get_dataset_names()
     for name in dataset_names:
         df = load_dataset(name)
         assert len(df) > 0
+
+
+def test_load_db() -> None:
+    dataset_names = get_db_names()
+    for name in dataset_names:
+        db = load_db(name)