Merge pull request #109 from CDCgov/dmb/catalog-generated-type-stubs

Generate precise catalog type stubs

Author: damonbayer

README.md

@@ -33,6 +33,12 @@ 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
+   ```
+   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

@@ -9,7 +9,7 @@
 from io import BytesIO
 from pathlib import PurePosixPath
 from types import SimpleNamespace
-from typing import TYPE_CHECKING, Any, Literal, overload
+from typing import Any, Literal, overload
 
 import pandas as pd
 import polars as pl
@@ -85,38 +85,14 @@ def get_all_catalogs() -> list:
 
 
 class CatalogNamespace(SimpleNamespace):
-    """Recursive namespace wrapper for statically-typed catalog access.
-
-    Runtime instances still behave like ``SimpleNamespace`` objects, but the
-    declared attributes let editors resolve endpoint methods on dynamic access
-    chains such as ``datacat.public.team.dataset.load``.
-    """
-
-    if TYPE_CHECKING:
-        load: "BlobEndpoint"
-        extract: "BlobEndpoint"
-        data: "BlobEndpoint"
-        _ledger_endpoint: "BlobEndpoint"
-        config: dict[str, Any]
-        __namespace_list__: list[str]
-
-    def __getattr__(self, name: str) -> "CatalogNamespace":
-        raise AttributeError(
-            f"{type(self).__name__!s} object has no attribute {name!r}"
-        )
+    """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
     blob endpoints for each 'stage' of the config (e.g., extract, load, stage_01)."""
 
-    if TYPE_CHECKING:
-        config: dict[str, Any]
-        load: "BlobEndpoint"
-        extract: "BlobEndpoint"
-        data: "BlobEndpoint"
-
     def __init__(self, config_path: str, defaults: dict, ns: str):
         """Basic functionality to interact with datasets to be included
         via the datasets configs.

cfa/dataops/stub_templates/__init__.py

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

cfa/dataops/stub_templates/catalog.pyi.mako

@@ -0,0 +1,123 @@
+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/init.pyi.mako

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