refactor: generics now used more carefully. Exactly where they are should be

Author: LeonidElkin

src/pysatl_core/distributions/computation.py

@@ -145,3 +145,6 @@ def fit(self, distribution: Distribution, **options: Any) -> FittedComputationMe
             Fitted method ready for evaluation.
         """
         return self.fitter(distribution, **options)
+
+
+type Method[In, Out] = AnalyticalComputation[In, Out] | FittedComputationMethod[In, Out]

src/pysatl_core/distributions/distribution.py

@@ -19,10 +19,9 @@
     from collections.abc import Mapping
     from typing import Any
 
-    from pysatl_core.distributions.computation import AnalyticalComputation
+    from pysatl_core.distributions.computation import AnalyticalComputation, Method
     from pysatl_core.distributions.strategies import (
         ComputationStrategy,
-        Method,
         SamplingStrategy,
     )
     from pysatl_core.distributions.support import Support
@@ -67,7 +66,7 @@ def analytical_computations(
     def sampling_strategy(self) -> SamplingStrategy: ...
 
     @property
-    def computation_strategy(self) -> ComputationStrategy[Any, Any]: ...
+    def computation_strategy(self) -> ComputationStrategy: ...
 
     @property
     def support(self) -> Support | None: ...

src/pysatl_core/distributions/fitters.py

@@ -27,7 +27,9 @@
     from typing import Any
 
     from pysatl_core.distributions.distribution import Distribution
-    from pysatl_core.types import GenericCharacteristicName, ScalarFunc
+    from pysatl_core.types import GenericCharacteristicName
+
+    type ScalarFunc = Callable[[float], float]
 
 
 def _resolve(distribution: Distribution, name: GenericCharacteristicName) -> ScalarFunc:

src/pysatl_core/distributions/strategies.py

@@ -21,14 +21,12 @@
 if TYPE_CHECKING:
     from typing import Any
 
-    from pysatl_core.distributions.computation import AnalyticalComputation, FittedComputationMethod
+    from pysatl_core.distributions.computation import FittedComputationMethod, Method
     from pysatl_core.distributions.distribution import Distribution
     from pysatl_core.types import GenericCharacteristicName
 
-type Method[In, Out] = AnalyticalComputation[In, Out] | FittedComputationMethod[In, Out]
 
-
-class ComputationStrategy[In, Out](Protocol):
+class ComputationStrategy(Protocol):
     """
     Protocol for strategies that resolve computation methods for characteristics.
 
@@ -42,10 +40,10 @@ class ComputationStrategy[In, Out](Protocol):
 
     def query_method(
         self, state: GenericCharacteristicName, distr: Distribution, **options: Any
-    ) -> Method[In, Out]: ...
+    ) -> Method[Any, Any]: ...
 
 
-class DefaultComputationStrategy[In, Out]:
+class DefaultComputationStrategy:
     """
     Default strategy for resolving characteristic computation methods.
 
@@ -71,7 +69,7 @@ class DefaultComputationStrategy[In, Out]:
 
     def __init__(self, enable_caching: bool = False) -> None:
         self.enable_caching = enable_caching
-        self._cache: dict[GenericCharacteristicName, FittedComputationMethod[In, Out]] = {}
+        self._cache: dict[GenericCharacteristicName, FittedComputationMethod[Any, Any]] = {}
         self._resolving: dict[int, set[GenericCharacteristicName]] = {}
 
     def _push_guard(self, distr: Distribution, state: GenericCharacteristicName) -> None:
@@ -103,7 +101,7 @@ def _pop_guard(self, distr: Distribution, state: GenericCharacteristicName) -> N
 
     def query_method(
         self, state: GenericCharacteristicName, distr: Distribution, **options: Any
-    ) -> Method[In, Out]:
+    ) -> Method[Any, Any]:
         """
         Resolve a computation method for the target characteristic.
 
@@ -164,7 +162,7 @@ def query_method(
                     continue
 
                 # Fit each edge along the path
-                last_fitted: FittedComputationMethod[In, Out] | None = None
+                last_fitted: FittedComputationMethod[Any, Any] | None = None
                 for edge in path:
                     fitted = edge.fit(distr, **options)
                     if self.enable_caching:

src/pysatl_core/families/distribution.py

@@ -158,7 +158,7 @@ def sampling_strategy(self) -> SamplingStrategy:
         return self.family.sampling_strategy
 
     @property
-    def computation_strategy(self) -> ComputationStrategy[Any, Any]:
+    def computation_strategy(self) -> ComputationStrategy:
         """Get the computation strategy for this distribution."""
         return self.family.computation_strategy
 

src/pysatl_core/families/parametric_family.py

@@ -41,15 +41,19 @@
     type SupportArg = Callable[[Parametrization], Support | None] | None
     type SupportResolver = Callable[[Parametrization], Support | None]
     type CharacteristicProvider = (
-        Mapping[ParametrizationName, CharacteristicFunction] | CharacteristicFunction
+        Mapping[ParametrizationName, CharacteristicFunction[Any, Any]]
+        | CharacteristicFunction[Any, Any]
     )
     type CharacteristicsMap = Mapping[GenericCharacteristicName, CharacteristicProvider]
+    type NonParametrizedCharacteristic[In, Out] = Callable[[], Out]
+    type CharacteristicFunction[In, Out] = (
+        NonParametrizedCharacteristic[In, Out] | ParametrizedCharacteristic[In, Out]
+    )
+
 
-type NonParametrizedCharacteristic = Callable[[], Any]
-type ParametrizedCharacteristic = (
-    Callable[[Parametrization, Any], Any] | Callable[[Parametrization], Any]
+type ParametrizedCharacteristic[In, Out] = (
+    Callable[[Parametrization, In], Out] | Callable[[Parametrization], Out]
 )
-type CharacteristicFunction = NonParametrizedCharacteristic | ParametrizedCharacteristic
 
 
 class ParametricFamily:
@@ -94,7 +98,7 @@ def __init__(
         distr_parametrizations: list[ParametrizationName],
         distr_characteristics: CharacteristicsMap,
         sampling_strategy: SamplingStrategy | None = None,
-        computation_strategy: ComputationStrategy[Any, Any] | None = None,
+        computation_strategy: ComputationStrategy | None = None,
         support_by_parametrization: SupportArg = None,
     ):
         self._name = name
@@ -124,16 +128,17 @@ def __init__(
         )
 
         def _process_char_val(
-            value: Mapping[ParametrizationName, CharacteristicFunction] | CharacteristicFunction,
-        ) -> dict[ParametrizationName, CharacteristicFunction]:
+            value: Mapping[ParametrizationName, CharacteristicFunction[Any, Any]]
+            | CharacteristicFunction[Any, Any],
+        ) -> dict[ParametrizationName, CharacteristicFunction[Any, Any]]:
             return (
                 dict(value)
                 if isinstance(value, Mapping)
                 else {self.base_parametrization_name: value}
             )
 
         self.distr_characteristics: dict[
-            GenericCharacteristicName, dict[ParametrizationName, CharacteristicFunction]
+            GenericCharacteristicName, dict[ParametrizationName, CharacteristicFunction[Any, Any]]
         ] = {key: _process_char_val(val) for key, val in distr_characteristics.items()}
 
         # Precompute analytical plan
@@ -236,9 +241,9 @@ def to_base(self, parameters: Parametrization) -> Parametrization:
         return parameters.transform_to_base_parametrization()
 
     @staticmethod
-    def _bind_parametrization(
-        func: CharacteristicFunction, params_obj: Parametrization
-    ) -> ComputationFunc[Any, Any]:
+    def _bind_parametrization[In, Out](
+        func: CharacteristicFunction[In, Out], params_obj: Parametrization
+    ) -> ComputationFunc[In, Out]:
         """Bind ``params_obj`` to ``func`` only when ``func`` can accept positional arguments.
 
         This allows parametrization-independent analytical providers to be written without
@@ -258,8 +263,8 @@ def _bind_parametrization(
         )
 
         return cast(
-            ComputationFunc[Any, Any],
-            partial(cast(ParametrizedCharacteristic, func), params_obj)
+            ComputationFunc[In, Out],
+            partial(cast(ParametrizedCharacteristic[In, Out], func), params_obj)
             if accepts_first_positional
             else func,
         )

src/pysatl_core/types.py

@@ -89,24 +89,24 @@ class EuclideanDistributionType(DistributionType):
 
 
 UnivariateContinuous = EuclideanDistributionType(kind=Kind.CONTINUOUS, dimension=1)
-"""Type for univariate continuous distributions."""
+"""Predefined DistributionType object for univariate continuous distributions."""
 
 UnivariateDiscrete = EuclideanDistributionType(kind=Kind.DISCRETE, dimension=1)
-"""Type for univariate discrete distributions."""
+"""Predefined DistributionType object for univariate discrete distributions."""
 
-NumPyNumber = np.floating[Any] | np.integer[Any]
+type NumPyNumber = np.floating[Any] | np.integer[Any]
 """Type alias for NumPy numeric types."""
 
-Number = NumPyNumber | int | float
+type Number = NumPyNumber | int | float
 """Type alias for all numeric types."""
 
-NumericArray = NDArray[NumPyNumber]
+type NumericArray = NDArray[NumPyNumber]
 """Type alias for numeric arrays."""
 
-ComplexArray = NDArray[np.complexfloating[Any]]
+type ComplexArray = NDArray[np.complexfloating[Any]]
 """Type alias for complex arrays."""
 
-BoolArray = NDArray[np.bool_]
+type BoolArray = NDArray[np.bool_]
 """Type alias for boolean arrays."""
 
 
@@ -253,9 +253,6 @@ def shape(self) -> ContinuousSupportShape1D:
 ``**options`` dynamically.
 """
 
-ScalarFunc = Callable[[float], float]
-"""Type alias for scalar functions (float -> float)."""
-
 
 class CharacteristicName(StrEnum):
     """
@@ -301,7 +298,6 @@ class FamilyName(StrEnum):
     "ParametrizationName",
     "ComputationFunc",
     "DistributionType",
-    "ScalarFunc",
     "Interval1D",
     "ContinuousSupportShape1D",
     "BoolArray",

tests/unit/distributions/test_registry.py

@@ -44,7 +44,7 @@ def test_configuration_continuous_presence_and_connectivity(self) -> None:
         assert view.find_path(CharacteristicName.PPF, CharacteristicName.CDF) is not None
 
         # Strategy resolves and roundtrips: CDF(PPF(q)) ~ q
-        strategy = DefaultComputationStrategy[float, float](enable_caching=False)
+        strategy = DefaultComputationStrategy(enable_caching=False)
         ppf = strategy.query_method(CharacteristicName.PPF, distr)
         cdf = strategy.query_method(CharacteristicName.CDF, distr)
         qs = np.linspace(1e-6, 1.0 - 1e-6, 7)
@@ -76,7 +76,7 @@ def test_configuration_discrete_requires_support_then_ok(self) -> None:
         assert view.find_path(CharacteristicName.CDF, CharacteristicName.PPF) is not None
         assert view.find_path(CharacteristicName.PPF, CharacteristicName.CDF) is not None
 
-        strategy = DefaultComputationStrategy[float, float](enable_caching=False)
+        strategy = DefaultComputationStrategy(enable_caching=False)
         cdf = strategy.query_method(CharacteristicName.CDF, distr)
         assert cdf(0.0) == pytest.approx(0.2, abs=1e-10)
         assert cdf(1.0) == pytest.approx(0.7, abs=1e-10)

tests/utils/mocks.py

@@ -81,7 +81,7 @@ def sampling_strategy(self) -> SamplingStrategy:
         return DefaultSamplingUnivariateStrategy()
 
     @property
-    def computation_strategy(self) -> ComputationStrategy[Any, Any]:
+    def computation_strategy(self) -> ComputationStrategy:
         """Computation strategy instance."""
         return DefaultComputationStrategy()