Add polars datastructures to Vector and Matrix type definitions.

Quantco · Sep 5, 2024 · 95b5058 · 95b5058
1 parent 0462b5f
commit 95b5058
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/metalearners/_typing.py b/metalearners/_typing.py
@@ -6,6 +6,7 @@
 
 import numpy as np
 import pandas as pd
+import polars as pl
 import scipy.sparse as sps
 
 PredictMethod = Literal["predict", "predict_proba"]
@@ -21,8 +22,8 @@
 Features = Collection[str] | Collection[int] | None
 
 # ruff is not happy about the usage of Union.
-Vector = Union[pd.Series, np.ndarray]  # noqa
-Matrix = Union[pd.DataFrame, np.ndarray, sps.csr_matrix]  # noqa
+Vector = Union[pd.Series, np.ndarray, pl.Series]  # noqa
+Matrix = Union[pd.DataFrame, np.ndarray, sps.csr_matrix, pl.DataFrame]  # noqa
 
 
 class _ScikitModel(Protocol):