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

docs/backcompat.md

@@ -103,6 +103,7 @@ Here are exceptions to our backwards compatibility policy:
   need to rethink Narwhals. However, we expect such radical changes to be exceedingly unlikely.
 - We may consider making some type hints more precise.
 - Anything labelled "unstable".
+- We may sometimes need to bump the minimum versions of supported backends.
 
 In general, decision are driven by use-cases, and we conduct a search of public GitHub repositories
 before making any change.
@@ -113,6 +114,11 @@ before making any change.
 
 The following are differences between the main Narwhals namespace and `narwhals.stable.v1`:
 
+- Since Narwhals 1.35:
+
+  - pandas' ordered categoricals get mapped to `nw.Enum` instead of `nw.Categorical`.
+  - `nw.Enum` must be provided `categories` at instantiation.
+
 - Since Narwhals 1.29.0, `LazyFrame.gather_every` has been deprecated from the main namespace.
 
 - Since Narwhals 1.24.1, an empty or all-null object-dtype pandas Series is inferred to

narwhals/_dask/utils.py

@@ -9,6 +9,7 @@
 from narwhals.dependencies import get_pandas
 from narwhals.dependencies import get_pyarrow
 from narwhals.utils import Implementation
+from narwhals.utils import Version
 from narwhals.utils import import_dtypes_module
 from narwhals.utils import isinstance_or_issubclass
 from narwhals.utils import parse_version
@@ -24,7 +25,6 @@
     from narwhals._dask.dataframe import DaskLazyFrame
     from narwhals._dask.expr import DaskExpr
     from narwhals.dtypes import DType
-    from narwhals.utils import Version
 
 
 def maybe_evaluate_expr(df: DaskLazyFrame, obj: DaskExpr | object) -> dx.Series | object:
@@ -125,6 +125,16 @@ def narwhals_to_native_dtype(dtype: DType | type[DType], version: Version) -> An
         return "object"  # pragma: no cover
     if isinstance_or_issubclass(dtype, dtypes.Boolean):
         return "bool"
+    if isinstance_or_issubclass(dtype, dtypes.Enum):
+        if version is Version.V1:
+            msg = "Converting to Enum is not supported in narwhals.stable.v1"
+            raise NotImplementedError(msg)
+
+        if dtype is dtypes.Enum:
+            msg = "Can not cast / initialize Enum without categories present"
+            raise ValueError(msg)
+        return get_pandas().CategoricalDtype(categories=dtype.categories, ordered=True)  # pyright: ignore[reportAttributeAccessIssue]
+
     if isinstance_or_issubclass(dtype, dtypes.Categorical):
         return "category"
     if isinstance_or_issubclass(dtype, dtypes.Datetime):

narwhals/_pandas_like/utils.py

@@ -211,7 +211,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()
@@ -249,7 +251,13 @@ 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 version is Version.V1:
+            return dtypes.Categorical()
+        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)
@@ -310,7 +318,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.
@@ -471,8 +479,19 @@ 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 version is Version.V1:
+            msg = "Converting to Enum is not supported in narwhals.stable.v1"
+            raise NotImplementedError(msg)
+
+        if dtype is dtypes.Enum:
+            msg = "Can not cast / initialize Enum without categories present"
+            raise ValueError(msg)
+
+        return implementation.to_native_namespace().CategoricalDtype(
+            categories=dtype.categories,  # pyright: ignore[reportAttributeAccessIssue]
+            ordered=True,
+        )
+
     if isinstance_or_issubclass(
         dtype, (dtypes.Struct, dtypes.Array, dtypes.List, dtypes.Time, dtypes.Binary)
     ):

narwhals/_polars/utils.py

@@ -17,6 +17,7 @@
 from narwhals.exceptions import InvalidOperationError
 from narwhals.exceptions import NarwhalsError
 from narwhals.exceptions import ShapeError
+from narwhals.utils import Version
 from narwhals.utils import import_dtypes_module
 from narwhals.utils import isinstance_or_issubclass
 
@@ -26,7 +27,6 @@
     from narwhals._polars.expr import PolarsExpr
     from narwhals._polars.series import PolarsSeries
     from narwhals.dtypes import DType
-    from narwhals.utils import Version
 
     T = TypeVar("T")
 
@@ -111,7 +111,9 @@ def native_to_narwhals_dtype(
     if dtype == pl.Categorical:
         return dtypes.Categorical()
     if dtype == pl.Enum:
-        return dtypes.Enum()
+        if version is Version.V1:
+            return dtypes.Enum()  # type: ignore[call-arg]
+        return dtypes.Enum(dtype.categories)  # type: ignore[attr-defined]
     if dtype == pl.Date:
         return dtypes.Date()
     if isinstance_or_issubclass(dtype, pl.Datetime):
@@ -186,8 +188,15 @@ 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 version is Version.V1:
+            msg = "Converting to Enum is not supported in narwhals.stable.v1"
+            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,6 +1,8 @@
 from __future__ import annotations
 
+import enum
 from collections import OrderedDict
+from collections import defaultdict
 from datetime import timezone
 from itertools import starmap
 from typing import TYPE_CHECKING
@@ -9,6 +11,7 @@
 from narwhals.utils import isinstance_or_issubclass
 
 if TYPE_CHECKING:
+    from typing import Iterable
     from typing import Iterator
     from typing import Sequence
 
@@ -443,9 +446,44 @@ 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[str] | type[enum.Enum]) -> None:
+        if isinstance(categories, type) and issubclass(categories, enum.Enum):
+            categories = (v.value for v in categories)
+
+        _categories, errors = [], defaultdict(set)
+        for cat in categories:
+            if cat in _categories:
+                errors["duplicates"].add(cat)
+
+            if not isinstance(cat, str):
+                errors["invalid"].add(cat)
+            _categories.append(cat)
+
+        if res := errors["invalid"]:
+            msg = f"{type(self).__name__} categories must be strings; found invalid: {res!r}"
+            raise TypeError(msg)
+
+        if res := errors["duplicates"]:
+            msg = f"{type(self).__name__} categories must be unique; found duplicate {res!r}"
+            raise ValueError(msg)
+
+        self.categories = tuple(_categories)
+
+    def __eq__(self: Self, other: object) -> bool:
+        # allow comparing object instances to class
+        if type(other) is type:
+            return other is Enum
+        return isinstance(other, type(self)) and self.categories == other.categories
+
+    def __hash__(self: Self) -> int:  # pragma: no cover
+        return hash((self.__class__, tuple(self.categories)))
+
+    def __repr__(self: Self) -> str:  # pragma: no cover
+        return f"{type(self).__name__}(categories={list(self.categories)!r})"
+
 
 class Field:
     """Definition of a single field within a `Struct` DataType.

narwhals/stable/v1/_dtypes.py

@@ -11,7 +11,7 @@
 from narwhals.dtypes import Decimal
 from narwhals.dtypes import DType
 from narwhals.dtypes import Duration as NwDuration
-from narwhals.dtypes import Enum
+from narwhals.dtypes import Enum as NwEnum
 from narwhals.dtypes import Field
 from narwhals.dtypes import Float32
 from narwhals.dtypes import Float64
@@ -72,6 +72,35 @@ def __hash__(self: Self) -> int:
         return hash(self.__class__)
 
 
+class Enum(NwEnum):
+    """A fixed categorical encoding of a unique set of strings.
+
+    Polars has an Enum data type, while pandas and PyArrow do not.
+
+    Examples:
+       >>> import polars as pl
+       >>> import narwhals.stable.v1 as nw
+       >>> data = ["beluga", "narwhal", "orca"]
+       >>> s_native = pl.Series(data, dtype=pl.Enum(data))
+       >>> nw.from_native(s_native, series_only=True).dtype
+       Enum
+    """
+
+    def __init__(self: Self) -> None:
+        super(NwEnum, self).__init__()
+
+    def __eq__(self, other: DType | type[DType]) -> bool:  # type: ignore[override]
+        if type(other) is type:
+            return other in {type(self), NwEnum}
+        return isinstance(other, type(self))
+
+    def __hash__(self: Self) -> int:  # pragma: no cover
+        return super(NwEnum, self).__hash__()
+
+    def __repr__(self: Self) -> str:  # pragma: no cover
+        return super(NwEnum, self).__repr__()
+
+
 __all__ = [
     "Array",
     "Binary",

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
@@ -400,3 +403,51 @@ def test_cast_decimal_to_native() -> None:
                 .with_columns(a=nw.col("a").cast(nw.Decimal()))
                 .to_native()
             )
+
+
+class FakeEnum(enum.Enum):
+    A = "A"
+    B = "B"
+
+
+@pytest.mark.parametrize(
+    "categories", [["a", "b"], [np.str_("a"), np.str_("b")], FakeEnum]
+)
+def test_enum_valid(categories: Iterable[Any] | type[enum.Enum]) -> None:
+    dtype = nw.Enum(categories)
+    assert dtype == nw.Enum
+    assert len(dtype.categories) == len([*categories])
+
+
+@pytest.mark.parametrize(
+    ("categories", "exception", "match"),
+    [
+        (["a", None], TypeError, "categories must be strings"),
+        (["a", float("nan")], TypeError, "categories must be strings"),
+        ([object()], TypeError, "categories must be strings"),
+        (enum.Enum("FakeEnum", "a b"), TypeError, "categories must be strings"),
+        (["a", "a"], ValueError, "categories must be unique"),
+    ],
+)
+def test_enum_errors(
+    categories: Iterable[Any], exception: type[Exception], match: str
+) -> None:
+    with pytest.raises(exception, match=match):
+        nw.Enum(categories)
+
+
+def test_enum_from_series() -> None:
+    pytest.importorskip("polars")
+    import polars as pl
+
+    elements = "a", "d", "e", "b", "c"
+    categories = pl.Series(elements)
+    categories_nw = nw.from_native(categories, series_only=True)
+    assert nw.Enum(categories_nw).categories == elements
+    assert nw.Enum(categories).categories == elements
+
+
+def test_enum_categories_immutable() -> None:
+    dtype = nw.Enum(["a", "b"])
+    with pytest.raises(TypeError, match="does not support item assignment"):
+        dtype.categories[0] = "c"  # type: ignore[index]

tests/frame/schema_test.py

@@ -176,6 +176,19 @@ def test_dtypes() -> None:
     assert df_from_pd.schema == df_from_pd.collect_schema() == expected
     assert {name: df_from_pd[name].dtype for name in df_from_pd.columns} == expected
 
+    df_from_pd = nw.from_native(df_pl.to_pandas(), eager_only=True)
+
+    pure_pd_expected = {
+        **expected,
+        "n": nw.Datetime,
+        "s": nw.Object,
+        "u": nw.Object,
+    }
+    assert df_from_pd.schema == df_from_pd.collect_schema() == pure_pd_expected
+    assert {
+        name: df_from_pd[name].dtype for name in df_from_pd.columns
+    } == pure_pd_expected
+
     df_from_pa = nw.from_native(df_pl.to_arrow(), eager_only=True)
 
     assert df_from_pa.schema == df_from_pa.collect_schema() == expected