Enable mypy.

Author: kklein

metalearners/_utils.py

@@ -13,6 +13,7 @@
 import polars as pl
 import scipy
 from narwhals.dependencies import is_into_dataframe, is_into_series
+from scipy.sparse import csr_matrix
 from sklearn.base import is_classifier, is_regressor
 from sklearn.ensemble import (
     HistGradientBoostingClassifier,
@@ -42,12 +43,13 @@ def index_matrix(matrix: Matrix, rows: Vector) -> Matrix:
         if not hasattr(rows, "to_numpy"):
             raise ValueError("rows couldn't be converted to numpy.")
         rows = rows.to_numpy()
+    if isinstance(matrix, np.ndarray) or isinstance(matrix, csr_matrix):
+        return matrix[rows, :]
     if is_into_dataframe(matrix):
-        matrix_nw = nw.from_native(matrix, eager_only=True)  # type: ignore
+        matrix_nw = nw.from_native(matrix, eager_only=True)
         if rows.dtype == "bool":
             return matrix_nw.filter(rows.tolist()).to_native()
         return matrix_nw[rows.tolist(), :].to_native()
-    return matrix[rows, :]
 
 
 def index_vector(vector: Vector, rows: Vector) -> Vector:
@@ -56,12 +58,14 @@ def index_vector(vector: Vector, rows: Vector) -> Vector:
         if not hasattr(rows, "to_numpy"):
             raise ValueError("rows couldn't be converted to numpy.")
         rows = rows.to_numpy()
+    if isinstance(vector, np.ndarray):
+        return vector[rows]
     if is_into_series(vector):
-        vector_nw = nw.from_native(vector, series_only=True, eager_only=True)  # type: ignore
+        vector_nw = nw.from_native(vector, series_only=True, eager_only=True)
         if rows.dtype == "bool":
             return vector_nw.filter(rows).to_native()
-        return vector_nw[rows].to_native()  # type: ignore
-    return vector[rows]
+        return vector_nw[rows].to_native()
+    raise TypeError(f"Encountered unexpected type of vector: {type(vector)}.")
 
 
 def are_pd_indices_equal(*args: pd.DataFrame | pd.Series) -> bool:

metalearners/metalearner.py

@@ -124,14 +124,16 @@ def _filter_x_columns(X: Matrix, feature_set: Features) -> Matrix:
         return X
     if len(feature_set) == 0:
         return np.ones((safe_len(X), 1))
+    if isinstance(X, np.ndarray):
+        return X[:, np.array(feature_set)]
     if nw.dependencies.is_into_dataframe(X):
-        X_nw = nw.from_native(X, eager_only=True)  # type: ignore
+        X_nw = nw.from_native(X, eager_only=True)
         if all(map(lambda x: isinstance(x, int), feature_set)):
-            return X_nw.select(nw.nth(feature_set)).to_native()  # type: ignore
+            return X_nw.select(nw.nth(feature_set)).to_native()
         if all(map(lambda x: isinstance(x, str), feature_set)):
-            return X_nw.select(feature_set).to_native()  # type: ignore
+            return X_nw.select(feature_set).to_native()
         raise ValueError("features must either be all ints or all strings.")
-    return X[:, np.array(feature_set)]
+    raise TypeError(f"Unexpected type of matrix: {type(X)}.")
 
 
 def _validate_n_folds_synchronize(n_folds: dict[str, int]) -> None: