python create Categorical Series without casting

pola-rs · Nov 28, 2021 · cea0f77 · cea0f77
1 parent 8f198f4
commit cea0f77
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 1 deletion.
diff --git a/py-polars/polars/datatypes_constructor.py b/py-polars/polars/datatypes_constructor.py
@@ -4,6 +4,7 @@
 
 from polars.datatypes import (
     Boolean,
+    Categorical,
     DataType,
     Date,
     Datetime,
@@ -46,6 +47,7 @@
         Boolean: PySeries.new_opt_bool,
         Utf8: PySeries.new_str,
         Object: PySeries.new_object,
+        Categorical: PySeries.new_str,
     }
 
 

diff --git a/py-polars/polars/internals/construction.py b/py-polars/polars/internals/construction.py
@@ -6,10 +6,12 @@
 
 from polars import internals as pli
 from polars.datatypes import (
+    Categorical,
     DataType,
     Date,
     Datetime,
     Float32,
+    Time,
     py_type_to_arrow_type,
     py_type_to_dtype,
 )
@@ -112,6 +114,11 @@ def sequence_to_pyseries(
             pyseries = pyseries.cast(str(Date), True)
         elif dtype == Datetime:
             pyseries = pyseries.cast(str(Datetime), True)
+        elif dtype == Time:
+            pyseries = pyseries.cast(str(Time), True)
+        elif dtype == Categorical:
+            pyseries = pyseries.cast(str(Categorical), True)
+
         return pyseries
 
     else:

diff --git a/py-polars/polars/testing.py b/py-polars/polars/testing.py
@@ -41,7 +41,11 @@ def assert_series_equal(
     atol: float = 1.0e-8,
 ) -> None:
     obj = "Series"
-    can_be_subtracted = hasattr(dtype_to_py_type(left.dtype), "__sub__")
+    try:
+        can_be_subtracted = hasattr(dtype_to_py_type(left.dtype), "__sub__")
+    except NotImplementedError:
+        can_be_subtracted = False
+
     check_exact = check_exact or not can_be_subtracted or left.dtype == Boolean
     if not (isinstance(left, Series) and isinstance(right, Series)):
         raise_assert_detail(obj, "Type mismatch", type(left), type(right))

diff --git a/py-polars/tests/test_series.py b/py-polars/tests/test_series.py
@@ -1057,3 +1057,10 @@ def test_reshape() -> None:
     # test lazy_dispatch
     out = pl.select(pl.lit(s).reshape((-1, 1))).to_series()
     assert out.series_equal(expected)
+
+
+def test_init_categorical() -> None:
+    for values in [[None], ["foo", "bar"], [None, "foo", "bar"]]:
+        expected = pl.Series("a", values, dtype=pl.Utf8).cast(pl.Categorical)
+        a = pl.Series("a", values, dtype=pl.Categorical)
+        testing.assert_series_equal(a, expected)