Merge pull request #137 from e10v/dev

feat: improve type annotations

Author: e10v

src/_strip_doctest_artifacts/__init__.py

@@ -3,22 +3,17 @@
 from __future__ import annotations
 
 import re
-from typing import TYPE_CHECKING
 
 import markdown
 import markdown.extensions
 import markdown.preprocessors
 
 
-if TYPE_CHECKING:
-    from typing import Any
-
-
 class StripDoctestArtifactsPreprocessor(markdown.preprocessors.Preprocessor):
     """A preprocessor that removes doctest artifacts."""
 
     def run(self, lines: list[str]) -> list[str]:
-        """Run th preprocessor."""
+        """Run the preprocessor."""
         return [_strip(line) for line in lines]
 
 
@@ -39,6 +34,6 @@ def extendMarkdown(self, md: markdown.Markdown) -> None:
         )
 
 
-def makeExtension(**kwargs: dict[str, Any]) -> StripDoctestArtifactsExtension:
+def makeExtension(**kwargs: dict[str, object]) -> StripDoctestArtifactsExtension:
     """A factory function for the extension, required by Python-Markdown."""
     return StripDoctestArtifactsExtension(**kwargs)

src/tea_tasting/aggr.py

@@ -14,7 +14,6 @@
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
-    from typing import Any
 
     import narwhals.typing  # noqa: TC004
 
@@ -238,7 +237,7 @@ def read_aggregates(
     mean_cols: Sequence[str],
     var_cols: Sequence[str],
     cov_cols: Sequence[tuple[str, str]],
-) -> dict[Any, Aggregates]:
+) -> dict[object, Aggregates]:
     ...
 
 @overload
@@ -261,7 +260,7 @@ def read_aggregates(
     mean_cols: Sequence[str],
     var_cols: Sequence[str],
     cov_cols: Sequence[tuple[str, str]],
-) -> dict[Any, Aggregates] | Aggregates:
+) -> dict[object, Aggregates] | Aggregates:
     """Extract aggregated statistics.
 
     Args:
@@ -346,7 +345,7 @@ def _read_aggr_ibis(
     mean_cols: Sequence[str],
     var_cols: Sequence[str],
     cov_cols: Sequence[tuple[str, str]],
-) -> list[dict[str, Any]]:
+) -> list[dict[str, int | float]]:
     covar_cols = tuple({*var_cols, *itertools.chain(*cov_cols)})
     backend = ibis.get_backend(data)
     var_op = ibis.expr.operations.Variance
@@ -402,7 +401,7 @@ def _read_aggr_narwhals(
     mean_cols: Sequence[str],
     var_cols: Sequence[str],
     cov_cols: Sequence[tuple[str, str]],
-) -> list[dict[str, Any]]:
+) -> list[dict[str, int | float]]:
     data = nw.from_native(data)
     if not isinstance(data, nw.LazyFrame):
         data = data.lazy()
@@ -464,15 +463,15 @@ def _demean_nw_col(col: str, group_col: str | None) -> nw.Expr:
 
 
 def _get_aggregates(
-    data: dict[str, Any],
+    data: dict[str, float | int],
     *,
     has_count: bool,
     mean_cols: Sequence[str],
     var_cols: Sequence[str],
     cov_cols: Sequence[tuple[str, str]],
 ) -> Aggregates:
     return Aggregates(
-        count_=data[_COUNT] if has_count else None,
+        count_=data[_COUNT] if has_count else None,  # type: ignore
         mean_={col: data[_MEAN.format(col)] for col in mean_cols},
         var_={col: data[_VAR.format(col)] for col in var_cols},
         cov_={cols: data[_COV.format(*cols)] for cols in cov_cols},

src/tea_tasting/config.py

@@ -4,17 +4,17 @@
 from __future__ import annotations
 
 import contextlib
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
 
 import tea_tasting.utils
 
 
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
-    from typing import Any, Literal
+    from typing import Literal
 
 
-_global_config = {
+_global_config: dict[str, object] = {
     "alpha": 0.05,
     "alternative": "two-sided",
     "confidence_level": 0.95,
@@ -27,7 +27,51 @@
 }
 
 
-def get_config(option: str | None = None) -> Any:
+@overload
+def get_config(option: Literal["alpha"]) -> float:
+    ...
+
+@overload
+def get_config(option: Literal["alternative"]) -> str:
+    ...
+
+@overload
+def get_config(option: Literal["confidence_level"]) -> float:
+    ...
+
+@overload
+def get_config(option: Literal["equal_var"]) -> bool:
+    ...
+
+@overload
+def get_config(option: Literal["n_obs"]) -> int | Sequence[int] | None:
+    ...
+
+@overload
+def get_config(option: Literal["n_resamples"]) -> str:
+    ...
+
+@overload
+def get_config(option: Literal["power"]) -> float:
+    ...
+
+@overload
+def get_config(option: Literal["ratio"]) -> float | int:
+    ...
+
+@overload
+def get_config(option: Literal["use_t"]) -> bool:
+    ...
+
+@overload
+def get_config(option: str) -> object:
+    ...
+
+@overload
+def get_config(option: None = None) -> dict[str, object]:
+    ...
+
+def get_config(option: str | None = None) -> object:
     """Retrieve the current settings of the global configuration.
 
     Args:
@@ -62,7 +106,7 @@ def set_config(
     power: float | None = None,
     ratio: float | int | None = None,
     use_t: bool | None = None,
-    **kwargs: Any,
+    **kwargs: object,
 ) -> None:
     """Update the global configuration with specified settings.
 
@@ -129,8 +173,8 @@ def config_context(
     power: float | None = None,
     ratio: float | int | None = None,
     use_t: bool | None = None,
-    **kwargs: Any,
-) -> Iterator[Any]:
+    **kwargs: object,
+) -> Iterator[object]:
     """A context manager that temporarily modifies the global configuration.
 
     Args:

src/tea_tasting/datasets.py

@@ -12,19 +12,22 @@
 
 
 if TYPE_CHECKING:
-    from typing import Any, Literal
+    from typing import Literal, TypeAlias
 
     import numpy.typing as npt
 
     try:
-        from pandas import DataFrame as PandasDataFrame
+        from pandas import DataFrame as _PandasDataFrame
     except ImportError:
-        from typing import Any as PandasDataFrame
+        _PandasDataFrame = object
 
     try:
-        from polars import DataFrame as PolarsDataFrame
+        from polars import DataFrame as _PolarsDataFrame
     except ImportError:
-        from typing import Any as PolarsDataFrame
+        _PolarsDataFrame = object
+
+    PandasDataFrame: TypeAlias = _PandasDataFrame  # type: ignore
+    PolarsDataFrame: TypeAlias = _PolarsDataFrame  # type: ignore
 
 
 @overload
@@ -610,9 +613,9 @@ def _check_params(
 
 
 def _avg_by_groups(
-    values: npt.NDArray[np.number[Any]],
-    groups: npt.NDArray[np.number[Any]],
-) -> npt.NDArray[np.number[Any]]:
+    values: npt.NDArray[np.number],
+    groups: npt.NDArray[np.number],
+) -> npt.NDArray[np.number]:
     return np.concatenate([
         np.full(v.shape, v.mean())
         for v in np.split(values, np.unique(groups, return_index=True)[1])

src/tea_tasting/experiment.py

@@ -34,7 +34,7 @@ class ExperimentResult(
         "pvalue",
     )
 
-    def to_dicts(self) -> tuple[dict[str, Any], ...]:
+    def to_dicts(self) -> tuple[dict[str, object], ...]:
         """Convert the result to a sequence of dictionaries.
 
         Examples:
@@ -77,11 +77,11 @@ def to_dicts(self) -> tuple[dict[str, Any], ...]:
         return tuple(
             {"metric": k} | (v if isinstance(v, dict) else v._asdict())
             for k, v in self.items()
-        )
+        )  # type: ignore
 
 
 class ExperimentResults(
-    UserDict[tuple[Any, Any], ExperimentResult],
+    UserDict[tuple[object, object], ExperimentResult],
     tea_tasting.utils.DictsReprMixin,
 ):
     """Experiment results for multiple pairs of variants."""
@@ -95,7 +95,7 @@ class ExperimentResults(
         "pvalue",
     )
 
-    def to_dicts(self) -> tuple[dict[str, Any], ...]:
+    def to_dicts(self) -> tuple[dict[str, object], ...]:
         """Convert the result to a sequence of dictionaries."""
         return tuple(
             {"variants": str(variants)} | metric_result
@@ -111,7 +111,7 @@ class ExperimentPowerResult(
     """Result of the analysis of power in a experiment."""
     default_keys = ("metric", "power", "effect_size", "rel_effect_size", "n_obs")
 
-    def to_dicts(self) -> tuple[dict[str, Any], ...]:
+    def to_dicts(self) -> tuple[dict[str, object], ...]:
         """Convert the result to a sequence of dictionaries."""
         dicts = ()
         for metric, results in self.items():
@@ -230,7 +230,7 @@ def __init__(
     def analyze(
         self,
         data: narwhals.typing.IntoFrame | ibis.expr.types.Table,
-        control: Any = None,
+        control: object = None,
         *,
         all_variants: Literal[False] = False,
     ) -> ExperimentResult:
@@ -240,7 +240,7 @@ def analyze(
     def analyze(
         self,
         data: narwhals.typing.IntoFrame | ibis.expr.types.Table,
-        control: Any = None,
+        control: object = None,
         *,
         all_variants: Literal[True] = True,
     ) -> ExperimentResults:
@@ -249,7 +249,7 @@ def analyze(
     def analyze(
         self,
         data: narwhals.typing.IntoFrame | ibis.expr.types.Table,
-        control: Any = None,
+        control: object = None,
         *,
         all_variants: bool = False,
     ) -> ExperimentResult | ExperimentResults:
@@ -273,7 +273,7 @@ def analyze(
             variants = granular_data.keys()
         else:
             variants = self._read_variants(data)
-        variants = sorted(variants)
+        variants = sorted(variants)  # type: ignore
 
         if control is not None:
             variant_pairs = tuple(
@@ -318,10 +318,10 @@ def _analyze_metric(
         self,
         metric: tea_tasting.metrics.MetricBase[Any],
         data: narwhals.typing.IntoFrame | ibis.expr.types.Table,
-        aggregated_data: dict[Any, tea_tasting.aggr.Aggregates] | None,
-        granular_data: dict[Any, pa.Table] | None,
-        control: Any,
-        treatment: Any,
+        aggregated_data: dict[object, tea_tasting.aggr.Aggregates] | None,
+        granular_data: dict[object, pa.Table] | None,
+        control: object,
+        treatment: object,
     ) -> tea_tasting.metrics.MetricResult:
         if (
             isinstance(metric, tea_tasting.metrics.MetricBaseAggregated)
@@ -342,8 +342,8 @@ def _read_data(
         self,
         data: narwhals.typing.IntoFrame | ibis.expr.types.Table,
     ) -> tuple[
-        dict[Any, tea_tasting.aggr.Aggregates] | None,
-        dict[Any, pa.Table] | None,
+        dict[object, tea_tasting.aggr.Aggregates] | None,
+        dict[object, pa.Table] | None,
     ]:
         aggr_cols = tea_tasting.metrics.AggrCols()
         gran_cols = set()
@@ -371,7 +371,7 @@ def _read_data(
     def _read_variants(
         self,
         data: narwhals.typing.IntoFrame | ibis.expr.types.Table,
-    ) -> list[Any]:
+    ) -> list[object]:
         if isinstance(data, ibis.expr.types.Table):
             return (
                 data.select(self.variant)