Merge pull request #98 from CDCgov/rr-68-type-hints

update to CatalogNamespace

Author: ryanraaschCDC

MANIFEST.in

@@ -1,3 +1,4 @@
 recursive-include cfa/dataops/reporting/reports
 recursive-include cfa/dataops/etl/transform_templates
 include cfa/dataops/config.ini
+include cfa/dataops/py.typed

README.md

@@ -33,6 +33,18 @@ This project provides data tools and low friction access to versioned datasets w
    df = datacat.{{dataset}}.load.get_dataframe()
    df.glimpse()
    ```
+6. Generate project-local type stubs for the installed catalog:
+   ```bash
+   dataops_catalog_stubs
+   ```
+
+   If using uv, run:
+   ```bash
+   uv run dataops_catalog_stubs
+   ```
+
+   This writes stubs under `typings/`, which Pyright/Pylance uses by default.
+   For mypy, include that directory with `MYPYPATH=typings`.
 
 Read the [Dataset User Guide](docs/data_user_guide.md) for more information about accessing datasets.
 

cfa/dataops/catalog.py

@@ -84,6 +84,10 @@ def get_all_catalogs() -> list:
 report_namespaces = get_dataset_dot_path(all_reports_ns_map)
 
 
+class CatalogNamespace(SimpleNamespace):
+    """Runtime namespace wrapper for catalog access."""
+
+
 class DatasetEndpoint:
     """The DatasetEndpoint class for including in the datacat namespace.
     This ends the namespace branching at a config file and creates all the
@@ -736,7 +740,7 @@ def save_dir_to_blob(
                 )
 
 
-def dict_to_sn(d: Any, defaults: dict = None, ns: str = "") -> SimpleNamespace:
+def dict_to_sn(d: Any, defaults: dict | None = None, ns: str = "") -> CatalogNamespace:
     """Simple recursive namespace construction
 
     Args:
@@ -745,9 +749,9 @@ def dict_to_sn(d: Any, defaults: dict = None, ns: str = "") -> SimpleNamespace:
         ns (str, optional): the current namespace path. Defaults to ''.
 
     Returns:
-        SimpleNamespace: namespace representation
+        CatalogNamespace: namespace representation
     """
-    x = SimpleNamespace()
+    x = CatalogNamespace()
     ns_prefix = f"{ns}." if ns != "" else ""
     _ = [
         setattr(
@@ -788,13 +792,13 @@ def dict_to_sn(d: Any, defaults: dict = None, ns: str = "") -> SimpleNamespace:
     rc.append(report_dict_to_sn({k: all_reports_ns_map[k]}))
 combined_reports_dict = {key: value for ns in rc for key, value in vars(ns).items()}
 
-datacat = SimpleNamespace(**combined_dict)
+datacat: CatalogNamespace = CatalogNamespace(**combined_dict)
 datacat.__setattr__("__namespace_list__", dataset_namespaces)
-reportcat = SimpleNamespace(**combined_reports_dict)
+reportcat: CatalogNamespace = CatalogNamespace(**combined_reports_dict)
 reportcat.__setattr__("__namespace_list__", report_namespaces)
 
 
-def _attach_schema_mock_functions(datacat: SimpleNamespace, catalogs: list) -> None:
+def _attach_schema_mock_functions(datacat: CatalogNamespace, catalogs: list) -> None:
     """Recursively walk the datacat namespace and attach mock_data functions to
     the extract and load BlobEndpoints of each DatasetEndpoint, sourced from
     a schema module co-located with the dataset.
@@ -816,12 +820,12 @@ def _attach_schema_mock_functions(datacat: SimpleNamespace, catalogs: list) -> N
         datacat.<catalog>.<team_path_segments>.<dataset>.load.mock_data()
 
     Args:
-        datacat (SimpleNamespace): the top-level datacat namespace
+        datacat (CatalogNamespace): the top-level datacat namespace
         catalogs (list): list of (catalog_namespace, catalog_name, catalog_path)
             tuples from get_all_catalogs()
     """
 
-    def _walk(ns: SimpleNamespace) -> None:
+    def _walk(ns: CatalogNamespace) -> None:
         for val in vars(ns).values():
             if isinstance(val, DatasetEndpoint):
                 # __ns_str__ is e.g. "public.stf.nhsn_hrd_prelim";
@@ -857,7 +861,7 @@ def _walk(ns: SimpleNamespace) -> None:
                                     "mock_data",
                                     getattr(mod, func_name),
                                 )
-            elif isinstance(val, SimpleNamespace):
+            elif isinstance(val, CatalogNamespace):
                 _walk(val)
 
     _walk(datacat)

cfa/dataops/py.typed

@@ -0,0 +1 @@
+"""Package resources for generated type stubs."""

cfa/dataops/stub_templates/__init__.py

@@ -0,0 +1,116 @@
+from collections.abc import Sequence
+from types import SimpleNamespace
+from typing import Any, Literal, overload
+
+import pandas as pd
+import polars as pl
+
+from .reporting.catalog import NotebookEndpoint
+
+def get_all_catalogs() -> list[tuple[str, str, str]]: ...
+
+class CatalogNamespace(SimpleNamespace):
+    pass
+
+class DatasetEndpoint:
+    config_path: str
+    defaults: dict[str, Any]
+    config: dict[str, Any]
+    __ns_str__: str
+    _ledger_location: dict[str, Any]
+
+class BlobEndpoint:
+    account: str
+    container: str
+    prefix: str
+    ledger_location: dict[str, Any]
+    is_ledger: bool
+    __ns_str__: str
+    def write_blob(
+        self,
+        file_buffer: bytes | Sequence[bytes],
+        path_after_prefix: str,
+        auto_version: bool = False,
+        append: bool = False,
+    ) -> None: ...
+    def read_blobs(
+        self,
+        version_spec: str | None = None,
+        selection: Literal["newest", "oldest", "all"] = "newest",
+        print_version: bool = True,
+    ) -> list[bytes]: ...
+    def read_csv(self, suffix: str) -> pd.DataFrame: ...
+    def get_versions(self) -> list[str]: ...
+    def get_file_ext(
+        self,
+        version_spec: str | None = None,
+        selection: Literal["newest", "oldest", "all"] = "newest",
+    ) -> str: ...
+    def download_version_to_local(
+        self,
+        local_path: str,
+        version_spec: str | None = None,
+        force: bool = False,
+        selection: Literal["newest", "oldest", "all"] = "newest",
+    ) -> bool: ...
+    @overload
+    def get_dataframe(
+        self,
+        output: Literal["pandas", "pd"] = "pandas",
+        version_spec: str | None = None,
+        selection: Literal["newest", "oldest"] = "newest",
+    ) -> pd.DataFrame: ...
+    @overload
+    def get_dataframe(
+        self,
+        output: Literal["polars", "pl"],
+        version_spec: str | None = None,
+        selection: Literal["newest", "oldest"] = "newest",
+    ) -> pl.DataFrame: ...
+    @overload
+    def get_dataframe(
+        self,
+        output: Literal["pl_lazy", "lazy"],
+        version_spec: str | None = None,
+        selection: Literal["newest", "oldest"] = "newest",
+    ) -> pl.LazyFrame: ...
+    def ledger_entry(self, action: str) -> None: ...
+    def save_dataframe(
+        self,
+        df: pd.DataFrame | pl.DataFrame,
+        path_after_prefix: str,
+        file_format: str = "parquet",
+        auto_version: bool = False,
+    ) -> None: ...
+    def save_file_to_blob(
+        self,
+        file_path: str,
+        path_after_prefix: str,
+        auto_version: bool = False,
+    ) -> None: ...
+    def save_dir_to_blob(
+        self,
+        dir_path: str,
+        path_after_prefix: str,
+        auto_version: bool = False,
+    ) -> None: ...
+
+def dict_to_sn(
+    d: Any,
+    defaults: dict[str, Any] | None = None,
+    ns: str = "",
+) -> CatalogNamespace: ...
+
+% for cls in model.classes:
+class ${cls.name}(${cls.base}):
+% if cls.attributes:
+% for attr in cls.attributes:
+    ${attr.name}: ${attr.type_name}
+% endfor
+% else:
+    pass
+% endif
+
+% endfor
+datacat: DataCatalog
+reportcat: ReportCatalog

cfa/dataops/stub_templates/catalog.pyi.mako

@@ -0,0 +1,4 @@
+from .catalog import datacat as datacat, reportcat as reportcat
+
+__version__: str
+__all__: list[str]