get dataset included in autodoc

matekadlicsko · Nov 25, 2022 · 856febe · 856febe
1 parent 988ac05
commit 856febe
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/causalpy/data/datasets.py b/causalpy/data/datasets.py
@@ -15,16 +15,20 @@
 }
 
 
-def get_data_home():
+def _get_data_home() -> pathlib.PosixPath:
     """Return the path of the data directory"""
     return pathlib.Path(cp.__file__).parents[1] / "causalpy" / "data"
 
 
-def load_data(dataset: str = None):
+def load_data(dataset: str = None) -> pd.DataFrame:
+    """Loads the requested dataset and returns a pandas DataFrame.
+
+    :param dataset: The desired dataset to load
+    """
 
     if dataset in DATASETS:
 
-        data_dir = get_data_home()
+        data_dir = _get_data_home()
         datafile = DATASETS[dataset]
         file_path = data_dir / datafile["filename"]
         return pd.read_csv(file_path)

diff --git a/docs/api_datasets.rst b/docs/api_datasets.rst
@@ -0,0 +1,9 @@
+:mod:`causalpy.data.datasets`
+=============================
+
+.. toctree::
+  :maxdepth: 1
+
+.. automodule:: causalpy.data.datasets
+   :members:
+   :undoc-members:
diff --git a/docs/index.rst b/docs/index.rst
@@ -113,6 +113,7 @@ Documentation outline
    api_pymc_experiments
    api_pymc_models
    api_plot_utils
+   api_datasets
 
 
 Index