refactor: distribution is abstract not the protocol

Author: LeonidElkin

src/pysatl_core/distributions/computation.py

@@ -175,12 +175,12 @@ def _bound(*args: Any, **kwargs: Any) -> Out:
         )
 
     @overload
-    def __call__(self, distribution: Distribution, **options: Any) -> Out: ...
+    def evaluate(self, distribution: Distribution, **options: Any) -> Out: ...
 
     @overload
-    def __call__(self, distribution: Distribution, data: In, **options: Any) -> Out: ...
+    def evaluate(self, distribution: Distribution, data: In, **options: Any) -> Out: ...
 
-    def __call__(self, distribution: Distribution, *args: Any, **options: Any) -> Out:
+    def evaluate(self, distribution: Distribution, *args: Any, **options: Any) -> Out:
         """Evaluate *direct* computation methods.
 
         This is only available for methods defined via ``evaluator``.
@@ -215,5 +215,15 @@ def fit(self, distribution: Distribution, **options: Any) -> FittedComputationMe
             )
         return self.fitter(distribution, **options)
 
+    @overload
+    def __call__(self, distribution: Distribution, **options: Any) -> Out: ...
+
+    @overload
+    def __call__(self, distribution: Distribution, data: In, **options: Any) -> Out: ...
+
+    def __call__(self, distribution: Distribution, *args: Any, **options: Any) -> Out:
+        """Fit if possible and then evaluate"""
+        return self.prepare(distribution, **options)(*args)
+
 
 type Method[In, Out] = AnalyticalComputation[In, Out] | FittedComputationMethod[In, Out]

src/pysatl_core/distributions/distribution.py

@@ -11,7 +11,8 @@
 __copyright__ = "Copyright (c) 2025 PySATL project"
 __license__ = "SPDX-License-Identifier: MIT"
 
-from typing import TYPE_CHECKING, Protocol, Self, cast, runtime_checkable
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Self, cast
 
 from pysatl_core.types import NumericArray
 
@@ -34,8 +35,7 @@
     )
 
 
-@runtime_checkable
-class Distribution(Protocol):
+class Distribution(ABC):
     """
     Protocol defining the interface for probability distributions.
 
@@ -58,22 +58,28 @@ class Distribution(Protocol):
     """
 
     @property
+    @abstractmethod
     def distribution_type(self) -> DistributionType: ...
 
     @property
+    @abstractmethod
     def analytical_computations(
         self,
     ) -> Mapping[GenericCharacteristicName, AnalyticalComputation[Any, Any]]: ...
 
     @property
+    @abstractmethod
     def sampling_strategy(self) -> SamplingStrategy: ...
 
     @property
+    @abstractmethod
     def computation_strategy(self) -> ComputationStrategy: ...
 
     @property
+    @abstractmethod
     def support(self) -> Support | None: ...
 
+    @abstractmethod
     def _clone_with_strategies(
         self,
         *,

src/pysatl_core/distributions/strategies.py

@@ -36,8 +36,6 @@ class ComputationStrategy(Protocol):
         Whether to cache fitted computation methods.
     """
 
-    enable_caching: bool
-
     def query_method(
         self, state: GenericCharacteristicName, distr: Distribution, **options: Any
     ) -> Method[Any, Any]: ...
@@ -59,7 +57,7 @@ class DefaultComputationStrategy:
 
     Attributes
     ----------
-    enable_caching : bool
+    _enable_caching : bool
         Whether caching is enabled.
     _cache : dict[str, FittedComputationMethod]
         Cache of fitted computation methods.
@@ -68,10 +66,14 @@ class DefaultComputationStrategy:
     """
 
     def __init__(self, enable_caching: bool = False) -> None:
-        self.enable_caching = enable_caching
+        self._enable_caching = enable_caching
         self._cache: dict[GenericCharacteristicName, FittedComputationMethod[Any, Any]] = {}
         self._resolving: dict[int, set[GenericCharacteristicName]] = {}
 
+    @property
+    def is_caching_enabled(self) -> bool:
+        return self._enable_caching
+
     def _push_guard(self, distr: Distribution, state: GenericCharacteristicName) -> None:
         """
         Push a characteristic onto the resolution stack to detect cycles.
@@ -135,7 +137,7 @@ def query_method(
             return distr.analytical_computations[state]
 
         # 2. Check cache if enabled
-        if self.enable_caching:
+        if self._enable_caching:
             cached = self._cache.get(state)
             if cached is not None:
                 return cached
@@ -165,7 +167,7 @@ def query_method(
                 last_fitted: FittedComputationMethod[Any, Any] | None = None
                 for edge in path:
                     fitted = edge.prepare(distr, **options)
-                    if self.enable_caching and edge.cacheable:
+                    if self._enable_caching and edge.cacheable:
                         self._cache[edge.target] = fitted
                     last_fitted = fitted
 

src/pysatl_core/families/parametrizations.py

@@ -14,9 +14,8 @@
 
 from abc import ABC
 from dataclasses import dataclass, is_dataclass
-from functools import wraps
 from inspect import isfunction
-from typing import TYPE_CHECKING, ParamSpec, dataclass_transform
+from typing import TYPE_CHECKING, dataclass_transform
 
 from pysatl_core.types import ParametrizationName
 
@@ -107,10 +106,7 @@ def transform_to_base_parametrization(self) -> Parametrization:
         return self
 
 
-P = ParamSpec("P")
-
-
-def constraint(description: str) -> Callable[[Callable[P, bool]], Callable[P, bool]]:
+def constraint[**P](description: str) -> Callable[[Callable[P, bool]], Callable[P, bool]]:
     """
     Decorator to mark an instance method as a parameter constraint.
 
@@ -133,13 +129,9 @@ def constraint(description: str) -> Callable[[Callable[P, bool]], Callable[P, bo
     """
 
     def decorator(func: Callable[P, bool]) -> Callable[P, bool]:
-        @wraps(func)
-        def wrapper(*args: P.args, **kwargs: P.kwargs) -> bool:
-            return func(*args, **kwargs)
-
-        setattr(wrapper, "__is_constraint", True)
-        setattr(wrapper, "__constraint_description", description)
-        return wrapper
+        setattr(func, "__is_constraint", True)
+        setattr(func, "__constraint_description", description)
+        return func
 
     return decorator