Fix type hints

Quantco · Aug 9, 2023 · f1ba304 · f1ba304
1 parent 5c064c2
commit f1ba304
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/tabmat/formula.py b/src/tabmat/formula.py
@@ -345,7 +345,7 @@ def to_tabmat(
         dtype: numpy.dtype = numpy.float64,
         sparse_threshold: float = 0.1,
         cat_threshold: int = 4,
-    ) -> DenseMatrix:
+    ) -> Union[SparseMatrix, DenseMatrix]:
         if (self.values != 0).mean() > sparse_threshold:
             return DenseMatrix(self.values)
         else:
@@ -439,7 +439,7 @@ def to_tabmat(
         dtype: numpy.dtype = numpy.float64,
         sparse_threshold: float = 0.1,
         cat_threshold: int = 4,
-    ) -> Union[CategoricalMatrix, SplitMatrix]:
+    ) -> Union[DenseMatrix, CategoricalMatrix, SplitMatrix]:
         codes = self.codes.copy()
         categories = self.categories.copy()
         if -2 in self.codes: