get dataset included in autodoc

Author: drbenvincent

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)

docs/api_datasets.rst

@@ -0,0 +1,9 @@
+:mod:`causalpy.data.datasets`
+=============================
+
+.. toctree::
+  :maxdepth: 1
+
+.. automodule:: causalpy.data.datasets
+   :members:
+   :undoc-members:

docs/index.rst

@@ -113,6 +113,7 @@ Documentation outline
    api_pymc_experiments
    api_pymc_models
    api_plot_utils
+   api_datasets
 
 
 Index