Distributing typing information (#127) (#135)

* Typing set up

* Typing information - auto generated

* Update typing information - switched to inline

Author: justin13601

pyproject.toml

@@ -29,6 +29,12 @@ dependencies = [
     "meds == 0.3.3",
 ]
 
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.package-data]
+your_package = ["*.pyi", "py.typed"]
+
 [tool.setuptools_scm]
 
 [project.scripts]

src/aces/__main__.py

@@ -131,7 +131,7 @@ def get_and_validate_label_schema(df: pl.DataFrame) -> pa.Table:
 
 
 @hydra.main(version_base=None, config_path=str(config_yaml.parent.resolve()), config_name=config_yaml.stem)
-def main(cfg: DictConfig):
+def main(cfg: DictConfig) -> None:
     import os
     from datetime import datetime
     from pathlib import Path

src/aces/config.py

@@ -9,6 +9,7 @@
 from dataclasses import field
 from datetime import timedelta
 from pathlib import Path
+from typing import Any
 
 import networkx as nx
 import polars as pl
@@ -28,7 +29,7 @@
 
 @dataclasses.dataclass
 class PlainPredicateConfig:
-    code: str | dict
+    code: str | dict[str, Any]
     value_min: float | None = None
     value_max: float | None = None
     value_min_inclusive: bool | None = None
@@ -294,7 +295,7 @@ class DerivedPredicateConfig:
     expr: str
     static: bool = False
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if not self.expr:
             raise ValueError("Derived predicates must have a non-empty expression field.")
 
@@ -652,7 +653,7 @@ class WindowConfig:
     index_timestamp: str | None = None
 
     @classmethod
-    def _check_reference(cls, reference: str):
+    def _check_reference(cls, reference: str) -> None:
         """Checks to ensure referenced events are valid."""
         err_str = (
             "Window boundary reference must be either a valid alphanumeric/'_' string "
@@ -713,7 +714,7 @@ def _parse_boundary(cls, boundary: str) -> dict[str, str]:
             cls._check_reference(ref)
             return {"referenced": ref, "offset": None, "event_bound": None, "occurs_before": None}
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         # Parse the has constraints from the string representation to the tuple representation
         if self.has is not None:
             for each_constraint in self.has:
@@ -1132,7 +1133,11 @@ class TaskExtractorConfig:
     index_timestamp_window: str | None = None
 
     @classmethod
-    def load(cls, config_path: str | Path, predicates_path: str | Path = None) -> TaskExtractorConfig:
+    def load(
+        cls: TaskExtractorConfig,
+        config_path: str | Path,
+        predicates_path: str | Path = None,
+    ) -> TaskExtractorConfig:
         """Load a configuration file from the given path and return it as a dict.
 
         Args:
@@ -1420,7 +1425,7 @@ def load(cls, config_path: str | Path, predicates_path: str | Path = None) -> Ta
 
         return cls(predicates=predicate_objs, trigger=trigger, windows=windows)
 
-    def _initialize_predicates(self):
+    def _initialize_predicates(self) -> None:
         """Initialize the predicates tree from the configuration object and check validity.
 
         Raises:
@@ -1467,7 +1472,7 @@ def _initialize_predicates(self):
                 f"Graph: {nx.write_network_text(self._predicate_dag_graph)}"
             )
 
-    def _initialize_windows(self):
+    def _initialize_windows(self) -> None:
         """Initialize the windows tree from the configuration object and check validity.
 
         Raises:
@@ -1614,7 +1619,7 @@ def _initialize_windows(self):
 
         self.window_nodes = window_nodes
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         self._initialize_predicates()
         self._initialize_windows()
 
@@ -1627,12 +1632,12 @@ def predicates_DAG(self) -> nx.DiGraph:
         return self._predicate_dag_graph
 
     @property
-    def plain_predicates(self) -> dict[str:PlainPredicateConfig]:
+    def plain_predicates(self) -> dict[str, PlainPredicateConfig]:
         """Returns a dictionary of plain predicates in {name: code} format."""
         return {p: cfg for p, cfg in self.predicates.items() if cfg.is_plain}
 
     @property
-    def derived_predicates(self) -> OrderedDict[str:DerivedPredicateConfig]:
+    def derived_predicates(self) -> OrderedDict[str, DerivedPredicateConfig]:
         """Returns an ordered dictionary mapping derived predicates to their configs in a proper order."""
         return {
             p: self.predicates[p]

src/aces/expand_shards.py

@@ -71,7 +71,7 @@ def expand_shards(*shards: str) -> str:
     return ",".join(result)
 
 
-def main():
+def main() -> None:
     print(expand_shards(*sys.argv[1:]))
 
 

src/aces/types.py

@@ -5,7 +5,9 @@
 """
 
 import dataclasses
+from collections.abc import Iterator
 from datetime import timedelta
+from typing import Any
 
 import polars as pl
 
@@ -59,14 +61,14 @@ class TemporalWindowBounds:
     offset: timedelta | None = None
 
     # Needed to make it accessible like a tuple.
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Any]:
         return (getattr(self, field.name) for field in dataclasses.fields(self))
 
     # Needed to make it scriptable.
-    def __getitem__(self, key):
+    def __getitem__(self, key: int) -> Any:
         return tuple(getattr(self, field.name) for field in dataclasses.fields(self))[key]
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.offset is None:
             self.offset = timedelta(0)
 
@@ -205,7 +207,7 @@ class ToEventWindowBounds:
     right_inclusive: bool
     offset: timedelta | None = None
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.end_event == "":
             raise ValueError("The 'end_event' must be a non-empty string.")
 
@@ -226,11 +228,11 @@ def __post_init__(self):
             self.offset = timedelta(0)
 
     # Needed to make it accessible like a tuple.
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Any]:
         return (getattr(self, field.name) for field in dataclasses.fields(self))
 
     # Needed to make it scriptable.
-    def __getitem__(self, key):
+    def __getitem__(self, key: int) -> Any:
         return tuple(getattr(self, field.name) for field in dataclasses.fields(self))[key]
 
     @property

src/aces/utils.py

@@ -1,16 +1,17 @@
 import io
 import os
 import sys
+from collections.abc import Generator
 from contextlib import contextmanager
 from datetime import timedelta
 
 import hydra
-from bigtree import print_tree
+from bigtree import Node, print_tree
 from loguru import logger
 from pytimeparse import parse
 
 
-def parse_timedelta(time_str: str) -> timedelta:
+def parse_timedelta(time_str: str = None) -> timedelta:
     """Parse a time string and return a timedelta object.
 
     Using time expression parser: https://github.com/wroberts/pytimeparse
@@ -40,7 +41,7 @@ def parse_timedelta(time_str: str) -> timedelta:
 
 
 @contextmanager
-def capture_output():
+def capture_output() -> Generator[io.StringIO, None, None]:
     """A context manager to capture stdout output.
 
     This can eventually be eliminated if https://github.com/kayjan/bigtree/issues/285 is resolved.
@@ -60,14 +61,14 @@ def capture_output():
         sys.stdout = old_out  # Restore the original stdout
 
 
-def log_tree(node):
+def log_tree(node: Node) -> None:
     """Logs the tree structure using logging.info."""
     with capture_output() as captured:
         print_tree(node, style="const_bold")  # This will print to the captured StringIO instead of stdout
     logger.info("\n" + captured.getvalue())  # Log the captured output
 
 
-def hydra_loguru_init(filename) -> None:
+def hydra_loguru_init(filename: str) -> None:
     """Must be called from a hydra main!"""
     hydra_path = hydra.core.hydra_config.HydraConfig.get().runtime.output_dir
     logger.add(os.path.join(hydra_path, filename))