Added Array protocol

Author: lucasimi

src/tdamapper/_common.py

@@ -2,19 +2,41 @@
 This module provides common functionalities for internal use.
 """
 
+from __future__ import annotations
+
 import cProfile
 import io
 import pstats
 import warnings
+from typing import Any, Callable, Protocol
 
 import numpy as np
+from numpy.typing import NDArray
 
 warnings.filterwarnings("default", category=DeprecationWarning, module=r"^tdamapper\.")
 
 
-def deprecated(msg):
-    def deprecated_func(func):
-        def wrapper(*args, **kwargs):
+class Array(Protocol):
+
+    def __getitem__(self, index: int) -> Any:
+        """
+        Get an item from the array.
+        """
+
+    def __len__(self) -> int:
+        """
+        Get the length of the array.
+        """
+
+    def __setitem__(self, index: int, value: Any) -> None:
+        """
+        Set an item in the array.
+        """
+
+
+def deprecated(msg: str) -> Callable[..., Any]:
+    def deprecated_func(func: Callable[..., Any]) -> Callable[..., Any]:
+        def wrapper(*args: list[Any], **kwargs: dict[str, Any]) -> Any:
             warnings.warn(msg, DeprecationWarning, stacklevel=2)
             return func(*args, **kwargs)
 
@@ -23,17 +45,19 @@ def wrapper(*args, **kwargs):
     return deprecated_func
 
 
-def warn_user(msg):
+def warn_user(msg: str) -> None:
     warnings.warn(msg, UserWarning, stacklevel=2)
 
 
 class EstimatorMixin:
 
-    def _is_sparse(self, X):
+    def _is_sparse(self, X: Array) -> bool:
         # simple alternative use scipy.sparse.issparse
         return hasattr(X, "toarray")
 
-    def _validate_X_y(self, X, y):
+    def _validate_X_y(
+        self, X: Array, y: Array
+    ) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
         if self._is_sparse(X):
             raise ValueError("Sparse data not supported.")
 
@@ -70,8 +94,9 @@ def _validate_X_y(self, X, y):
 
         return X, y
 
-    def _set_n_features_in(self, X):
-        self.n_features_in_ = X.shape[1]
+    def _set_n_features_in(self, X: Array) -> None:
+        if hasattr(X, "shape"):
+            self.n_features_in_ = X.shape[1]
 
 
 class ParamsMixin:
@@ -80,16 +105,16 @@ class ParamsMixin:
     scikit-learn `get_params` and `set_params`.
     """
 
-    def _is_param_public(self, k):
+    def _is_param_public(self, k: str) -> bool:
         return (not k.startswith("_")) and (not k.endswith("_"))
 
-    def _split_param(self, k):
+    def _split_param(self, k: str) -> tuple[str, str]:
         k_split = k.split("__")
         outer = k_split[0]
         inner = "__".join(k_split[1:])
         return outer, inner
 
-    def get_params(self, deep=True):
+    def get_params(self, deep: bool = True) -> dict[str, Any]:
         """
         Get all public parameters of the object as a dictionary.
 
@@ -105,7 +130,7 @@ def get_params(self, deep=True):
                         params[f"{k}__{_k}"] = _v
         return params
 
-    def set_params(self, **params):
+    def set_params(self, **params: dict[str, Any]) -> ParamsMixin:
         """
         Set public parameters. Only updates attributes that already exist.
         """
@@ -124,7 +149,7 @@ def set_params(self, **params):
                 k_attr.set_params(**{k_inner: v})
         return self
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         obj_noargs = type(self)()
         args_repr = []
         for k, v in self.__dict__.items():
@@ -136,7 +161,7 @@ def __repr__(self):
         return f"{self.__class__.__name__}({', '.join(args_repr)})"
 
 
-def clone(obj):
+def clone(obj: Any) -> Any:
     """
     Clone an estimator, returning a new one, unfitted, having the same public
     parameters.
@@ -152,9 +177,9 @@ def clone(obj):
     return obj_noargs
 
 
-def profile(n_lines=10):
-    def decorator(func):
-        def wrapper(*args, **kwargs):
+def profile(n_lines: int = 10) -> Callable[..., Any]:
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        def wrapper(*args: list[Any], **kwargs: dict[str, Any]) -> Any:
             profiler = cProfile.Profile()
             profiler.enable()
             result = func(*args, **kwargs)