Support Enum

pandera/engines/polars_engine.py

@@ -680,11 +680,49 @@
 
     @classmethod
     def from_parametrized_dtype(cls, polars_dtype: pl.Categorical):
-        """Convert a :class:`polars.Decimal` to
-        a Pandera :class:`pandera.engines.polars_engine.Decimal`."""
+        """Convert a :class:`polars.Categorical` to
+        a Pandera :class:`pandera.engines.polars_engine.Categorical`."""
         return cls(ordering=polars_dtype.ordering)
 
 
+@Engine.register_dtype(equivalents=[pl.Enum])
+@immutable(init=True)
+class Enum(DataType):
+    """Polars enum data type."""
+
+    type = pl.Enum
+
+    categories: pl.Series
+
+    def __init__(  # pylint:disable=super-init-not-called
+        self,
+        categories: Union[pl.Series, Iterable[str], None] = None,
+    ) -> None:
+        object.__setattr__(self, "categories", categories)
+        object.__setattr__(self, "type", pl.Enum(categories=categories))
+
+    @classmethod
+    def from_parametrized_dtype(cls, polars_dtype: pl.Enum):
+        """Convert a :class:`polars.Enum` to
+        a Pandera :class:`pandera.engines.polars_engine.Enum`."""
+        return cls(categories=polars_dtype.categories)
+
+    def check(
+        self,
+        pandera_dtype: dtypes.DataType,
+        data_container: Optional[PolarsDataContainer] = None,
+    ) -> Union[bool, Iterable[bool]]:
+        try:
+            pandera_dtype = Engine.dtype(pandera_dtype)
+        except TypeError:
+            return False
+
+        return (
+            self.type == pandera_dtype.type
+            and (self.type.categories == pandera_dtype.categories).all()  # type: ignore
+        )
+
+
 @Engine.register_dtype(
     equivalents=["category", dtypes.Category, dtypes.Category()]
 )

tests/polars/test_polars_dtypes.py

@@ -409,6 +409,12 @@ def test_polars_struct_nested_type(inner_dtype_cls):
             pl.List(pl.Object()),
             pl.LazyFrame({"0": [[1.0, 2.0, 3.0]]}),
         ],
+        # Enum
+        [
+            pl.Enum(categories=["yes", "no"]),
+            pl.Enum(categories=["yes", "no", "?"]),
+            pl.LazyFrame({"0": ["yes", "yes", "no"]}),
+        ],
         # Struct
         [
             pl.Struct({"a": pl.Utf8(), "b": pl.Int64(), "c": pl.Float64()}),