feat: let nw.Enum accept categories, map pandas ordered categorical to Enum (only in main namespace, not stable.v1)

narwhals/_pandas_like/utils.py

@@ -314,7 +314,9 @@ def rename(
 
 
 @functools.lru_cache(maxsize=16)
-def non_object_native_to_narwhals_dtype(dtype: str, version: Version) -> DType:
+def non_object_native_to_narwhals_dtype(native_dtype: Any, version: Version) -> DType:
+    dtype = str(native_dtype)
+
     dtypes = import_dtypes_module(version)
     if dtype in {"int64", "Int64", "Int64[pyarrow]", "int64[pyarrow]"}:
         return dtypes.Int64()
@@ -352,7 +354,11 @@ def non_object_native_to_narwhals_dtype(dtype: str, version: Version) -> DType:
         return dtypes.String()
     if dtype in {"bool", "boolean", "boolean[pyarrow]", "bool[pyarrow]"}:
         return dtypes.Boolean()
-    if dtype == "category" or dtype.startswith("dictionary<"):
+    if dtype.startswith("dictionary<"):
+        return dtypes.Categorical()
+    if dtype == "category":
+        if native_dtype.ordered:
+            return dtypes.Enum(categories=native_dtype.categories)
         return dtypes.Categorical()
     if (match_ := PATTERN_PD_DATETIME.match(dtype)) or (
         match_ := PATTERN_PA_DATETIME.match(dtype)
@@ -413,7 +419,7 @@ def native_to_narwhals_dtype(
             return arrow_native_to_narwhals_dtype(native_dtype.to_arrow(), version)
         return arrow_native_to_narwhals_dtype(native_dtype.pyarrow_dtype, version)
     if str_dtype != "object":
-        return non_object_native_to_narwhals_dtype(str_dtype, version)
+        return non_object_native_to_narwhals_dtype(native_dtype, version)
     elif implementation is Implementation.DASK:
         # Per conversations with their maintainers, they don't support arbitrary
         # objects, so we can just return String.
@@ -569,8 +575,17 @@ def narwhals_to_native_dtype(  # noqa: PLR0915
             msg = "PyArrow>=11.0.0 is required for `Date` dtype."
         return "date32[pyarrow]"
     if isinstance_or_issubclass(dtype, dtypes.Enum):
-        msg = "Converting to Enum is not (yet) supported"
-        raise NotImplementedError(msg)
+        if isinstance(dtype, dtypes.Enum):
+            try:
+                import pandas as pd  # ignore-banned-import
+            except ImportError as exc:  # pragma: no cover
+                msg = f"Unable to convert to {dtype} to to the following exception: {exc.msg}"
+                raise ImportError(msg) from exc
+            return pd.CategoricalDtype(categories=dtype.categories, ordered=True)
+        else:
+            msg = "Can not cast / initialize Enum without categories present"
+            raise ValueError(msg)
+
     if isinstance_or_issubclass(
         dtype, (dtypes.Struct, dtypes.Array, dtypes.List, dtypes.Time, dtypes.Binary)
     ):

narwhals/_polars/utils.py

@@ -110,7 +110,7 @@ def native_to_narwhals_dtype(
     if dtype == pl.Categorical:
         return dtypes.Categorical()
     if dtype == pl.Enum:
-        return dtypes.Enum()
+        return dtypes.Enum(dtype.categories)  # type: ignore[attr-defined]
     if dtype == pl.Date:
         return dtypes.Date()
     if dtype == pl.Datetime:
@@ -185,8 +185,10 @@ def narwhals_to_native_dtype(
     if dtype == dtypes.Categorical:
         return pl.Categorical()
     if dtype == dtypes.Enum:
-        msg = "Converting to Enum is not (yet) supported"
-        raise NotImplementedError(msg)
+        if dtype is dtypes.Enum:
+            msg = "Can not cast / initialize Enum without categories present"
+            raise ValueError(msg)
+        return pl.Enum(dtype.categories)  # type: ignore[union-attr]
     if dtype == dtypes.Date:
         return pl.Date()
     if dtype == dtypes.Time:

narwhals/dtypes.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import enum
 from collections import OrderedDict
 from datetime import timezone
 from itertools import starmap
@@ -9,6 +10,8 @@
 from narwhals.utils import isinstance_or_issubclass
 
 if TYPE_CHECKING:
+    from collections.abc import Hashable
+    from typing import Iterable
     from typing import Iterator
     from typing import Sequence
 
@@ -443,9 +446,32 @@ class Enum(DType):
        >>> data = ["beluga", "narwhal", "orca"]
        >>> s_native = pl.Series(data, dtype=pl.Enum(data))
        >>> nw.from_native(s_native, series_only=True).dtype
-       Enum
+       Enum(categories=['beluga', 'narwhal', 'orca'])
     """
 
+    def __init__(self, categories: Iterable[Hashable] | type[enum.Enum]) -> None:
+        # TODO(Unassigned): pandas errors on NaN, NA, NaT OR duplicated value category
+        #       Polars errors on Null, NaN OR duplicated OR any non-string category
+        #       should the intersection of the above be caught at the narwhals layer?
+        if isinstance(categories, type) and issubclass(categories, enum.Enum):
+            categories = (getattr(v, "value", v) for v in categories.__members__.values())
+        self.categories = [*categories]
+
+    def __eq__(self: Self, other: object) -> bool:
+        # allow comparing object instances to class
+        if type(other) is type and issubclass(other, self.__class__):
+            return True
+        elif isinstance(other, self.__class__):
+            return self.categories == other.categories
+        return False
+
+    def __hash__(self: Self) -> int:  # pragma: no cover
+        return hash((self.__class__, tuple(self.categories)))
+
+    def __repr__(self: Self) -> str:  # pragma: no cover
+        class_name = self.__class__.__name__
+        return f"{class_name}(categories={self.categories!r})"
+
 
 class Field:
     """Definition of a single field within a `Struct` DataType.

tests/dtypes_test.py

@@ -1,9 +1,12 @@
 from __future__ import annotations
 
+import enum
 from datetime import datetime
 from datetime import timedelta
 from datetime import timezone
 from typing import TYPE_CHECKING
+from typing import Any
+from typing import Iterable
 from typing import Literal
 
 import numpy as np
@@ -393,3 +396,11 @@ def test_cast_decimal_to_native() -> None:
                 .with_columns(a=nw.col("a").cast(nw.Decimal()))
                 .to_native()
             )
+
+
+@pytest.mark.parametrize(
+    "categories", [["a", "b"], ["a", None], [1, 2], enum.Enum("test", "a b")]
+)
+def test_enum_valid(categories: Iterable[Any] | type[enum.Enum]) -> None:
+    dtype = nw.Enum(categories)
+    assert dtype == nw.Enum

tests/series_only/cast_test.py

@@ -117,25 +117,25 @@ def test_cast_to_enum_polars() -> None:
     pytest.importorskip("polars")
     import polars as pl
 
-    # we don't yet support metadata in dtypes, so for now disallow this
-    # seems like a very niche use case anyway, and allowing it later wouldn't be
-    # backwards-incompatible
     df_pl = pl.DataFrame({"a": ["a", "b"]}, schema={"a": pl.Categorical})
     with pytest.raises(
-        NotImplementedError, match=r"Converting to Enum is not \(yet\) supported"
+        ValueError, match="Can not cast / initialize Enum without categories present"
     ):
         nw.from_native(df_pl).select(nw.col("a").cast(nw.Enum))
 
+    df_nw_pl = nw.from_native(df_pl).select(nw.col("a").cast(nw.Enum(["a", "b"])))
+    assert df_nw_pl.schema == {"a": nw.Enum(["a", "b"])}
+
 
 def test_cast_to_enum_pandas() -> None:
     pytest.importorskip("pandas")
     import pandas as pd
 
-    # we don't yet support metadata in dtypes, so for now disallow this
-    # seems like a very niche use case anyway, and allowing it later wouldn't be
-    # backwards-incompatible
     df_pd = pd.DataFrame({"a": ["a", "b"]}, dtype="category")
     with pytest.raises(
-        NotImplementedError, match=r"Converting to Enum is not \(yet\) supported"
+        ValueError, match="Can not cast / initialize Enum without categories present"
     ):
         nw.from_native(df_pd).select(nw.col("a").cast(nw.Enum))
+
+    df_nw_pd = nw.from_native(df_pd).select(nw.col("a").cast(nw.Enum(["a", "b"])))
+    assert df_nw_pd.schema == {"a": nw.Enum(["a", "b"])}