Use IntEnum for ccflow.Enum with csp 0.19 or newer

ccflow/enums.py

@@ -1,6 +1,8 @@
 """We define our own enum type for use in ccflow.
 The implementation of this enum depends on whether csp is found in the environment or not.
 If it is found, it uses csp.Enum for seamless communication with the C++ layer of csp, otherwise it uses enum.Enum.
+Since csp 0.19, csp accepts any IntEnum natively, so for those versions it uses an IntEnum matching csp.Enum
+and does not import csp.
 We define a new subclass whose parent is either of the above, so that testing whether something is an enum
 (i.e. isinstance, or type hints) does not need to specify either of the implementations,
 i.e. isinstance(x, ccflow.enums.Enum).
@@ -28,36 +30,79 @@
 
 from enum import (
     Enum as BaseEnum,
+    IntEnum,
     auto,
 )
 
 BaseEnum.auto = auto
 
 _CSP_ENUM = False
+_CSP_DYNAMIC_ENUM = False
+
+
+def _csp_version() -> tuple[int, int] | None:
+    """Return the installed csp (major, minor) from package metadata, without importing csp."""
+    from importlib.metadata import PackageNotFoundError, version
 
-if not TYPE_CHECKING and environ.get("CCFLOW_NO_CSP", "0") != "1":
     try:
-        from csp import DynamicEnum, Enum as BaseEnum
+        major, minor = version("csp").split(".")[:2]
+        return int(major), int(minor)
+    except (PackageNotFoundError, ValueError):
+        return None
 
-        auto = BaseEnum.auto
 
-        _CSP_ENUM = True
+class _IntEnum(IntEnum):
+    """Matches csp.Enum from csp 0.19, which csp accepts natively along with any other IntEnum."""
+
+    @classmethod
+    def _missing_(cls, value):
+        # Allow construction by name, i.e. MyEnum("FOO") as well as MyEnum(1)
+        if isinstance(value, str):
+            return cls._member_map_.get(value)
+        return None
+
+    def __str__(self):
+        return f"{type(self).__name__}.{self.name}"
+
+    __hash__ = int.__hash__
+
+    def __eq__(self, other):
+        return type(self) is type(other) and self.value == other.value
 
-    except ImportError:
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+
+if not TYPE_CHECKING and environ.get("CCFLOW_NO_CSP", "0") != "1":
+    if (_csp_version() or (0, 0)) >= (0, 19):
+        # Since csp 0.19, csp treats any IntEnum as a native enum, so csp (and its import cost) is not needed here.
+        BaseEnum = _IntEnum
+        _CSP_ENUM = True
+    else:
         try:
-            from enum import auto
+            from csp import DynamicEnum, Enum as BaseEnum
 
-            from csp.impl.enum import DynamicEnum, Enum as BaseEnum
+            auto = BaseEnum.auto
 
             _CSP_ENUM = True
+            _CSP_DYNAMIC_ENUM = True
 
         except ImportError:
-            # if csp is not installed, rely on python Enum
-            pass
+            try:
+                from enum import auto
+
+                from csp.impl.enum import DynamicEnum, Enum as BaseEnum
+
+                _CSP_ENUM = True
+                _CSP_DYNAMIC_ENUM = True
+
+            except ImportError:
+                # if csp is not installed, rely on python Enum
+                pass
 
 
 class Enum(BaseEnum):
-    """Hybrid enum type based on csp.Enum if csp is available, otherwise enum.Enum"""
+    """Hybrid enum type based on csp.Enum (or an equivalent IntEnum for csp >= 0.19) if csp is available, otherwise enum.Enum"""
 
     # NOTE: we overload the default numbering
     # behavior to start at 0 and match to csp
@@ -116,7 +161,7 @@ def make_enum(name: str, values: dict[str, float] | list, start=0):
     """
 
     module_name = inspect.currentframe().f_back.f_globals["__name__"]
-    if _CSP_ENUM:
+    if _CSP_DYNAMIC_ENUM:
         e = DynamicEnum(name, values, start, module_name=module_name)
         # DynamicEnum is a BaseEnum (csp.Enum or csp.impl.enum.Enum) but it is not a ccflow.enums.Enum. To make
         # issubclass and isinstance return True when testing against ccflow.enums.Enum, we need to add it as a

ccflow/examples/narwhals_pipelines.ipynb

@@ -192,7 +192,7 @@
    ],
    "source": [
     "def q1_monolith(lineitem):\n",
-    "    cutoff = datetime(1998, 9, 2)  # noqa: DTZ001\n",
+    "    cutoff = datetime(1998, 9, 2)\n",
     "    return (\n",
     "        lineitem.filter(nw.col(\"l_shipdate\") <= cutoff)\n",
     "        .with_columns(\n",
@@ -370,7 +370,7 @@
     "        return df.filter(nw.col(\"l_shipdate\") <= self.cutoff)\n",
     "\n",
     "\n",
-    "filtered = lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2)))  # noqa: DTZ001\n",
+    "filtered = lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2)))\n",
     "filtered.collect().shape"
    ]
   },
@@ -434,7 +434,7 @@
     "        )\n",
     "\n",
     "\n",
-    "lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2))).pipe(DeriveLineItemMetrics()).collect().columns  # noqa: DTZ001"
+    "lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2))).pipe(DeriveLineItemMetrics()).collect().columns"
    ]
   },
   {
@@ -541,7 +541,7 @@
     }
    ],
    "source": [
-    "result = lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2))).pipe(DeriveLineItemMetrics()).pipe(SummarizeByReturnStatus())  # noqa: DTZ001\n",
+    "result = lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2))).pipe(DeriveLineItemMetrics()).pipe(SummarizeByReturnStatus())\n",
     "result.collect().to_native()"
    ]
   },
@@ -710,7 +710,7 @@
    "source": [
     "q1_preproc = SequenceTransform(\n",
     "    transforms=[\n",
-    "        FilterByShipDate(cutoff=datetime(1998, 9, 2)),  # noqa: DTZ001\n",
+    "        FilterByShipDate(cutoff=datetime(1998, 9, 2)),\n",
     "        DeriveLineItemMetrics(),\n",
     "        SummarizeByReturnStatus(),\n",
     "    ],\n",
@@ -860,7 +860,7 @@
     "q1_pipeline = NarwhalsPipelineModel(\n",
     "    source=registry[\"/table/lineitem\"],\n",
     "    transforms=[\n",
-    "        FilterByShipDate(cutoff=datetime(1998, 9, 2)),  # noqa: DTZ001\n",
+    "        FilterByShipDate(cutoff=datetime(1998, 9, 2)),\n",
     "        DeriveLineItemMetrics(),\n",
     "        SummarizeByReturnStatus(),\n",
     "    ],\n",
@@ -1133,7 +1133,7 @@
     "        )\n",
     "\n",
     "\n",
-    "cutoff = datetime(1995, 3, 15)  # noqa: DTZ001\n",
+    "cutoff = datetime(1995, 3, 15)\n",
     "\n",
     "q3_pipeline = NarwhalsPipelineModel(\n",
     "    source=registry[\"/table/customer\"],\n",

ccflow/tests/enums/test_enums.py

@@ -64,9 +64,12 @@ def test_init_parent_csp(self):
         mock_csp.Enum = TestEnumType
         mock_csp.DynamicEnum = MagicMock()
         mock_csp.EnumMeta = MagicMock()
-        with patch.dict(
-            "sys.modules",
-            {"csp": mock_csp, "csp.impl": mock_csp, "csp.impl.enum": mock_csp},
+        with (
+            patch.dict(
+                "sys.modules",
+                {"csp": mock_csp, "csp.impl": mock_csp, "csp.impl.enum": mock_csp},
+            ),
+            patch("importlib.metadata.version", return_value="0.18.0"),
         ):
             importlib.invalidate_caches()
 
@@ -80,6 +83,34 @@ def test_init_parent_csp(self):
 
             self.assertTrue(issubclass(Enum, csp.Enum))
 
+    def test_init_parent_csp_int_enum(self):
+        """Test initialization of ccflow.enum.Enum when csp >= 0.19 is installed, which accepts any IntEnum."""
+        # csp must not be imported on this path, so make importing it fail
+        with patch.dict("sys.modules", {"csp": None}), patch("importlib.metadata.version", return_value="0.19.0"):
+            import ccflow.enums
+
+            importlib.reload(ccflow.enums)
+
+            from ccflow.enums import Enum, auto, make_enum
+
+            self.assertTrue(issubclass(Enum, enum.IntEnum))
+            self.assertTrue(ccflow.enums._CSP_ENUM)
+
+            class MyEnum(Enum):
+                A = auto()
+                B = auto()
+
+            self.assertEqual(MyEnum.A.value, 0)
+            self.assertEqual(MyEnum("B"), MyEnum.B)
+            self.assertEqual(MyEnum(1), MyEnum.B)
+            self.assertEqual(str(MyEnum.A), "MyEnum.A")
+            self.assertNotEqual(MyEnum.A, 0)
+            self.assertEqual(hash(MyEnum.A), hash(0))
+
+            MyDynamicEnum = make_enum("MyDynamicEnum", ["A", "B"], start=2)
+            self.assertEqual(MyDynamicEnum.A.value, 2)
+            self.assertTrue(issubclass(MyDynamicEnum, Enum))
+
     def test_init(self):
         """Test standard enum behavior. This test should pass in all environments (both with and without csp).
         No other enum functionality should be depended-upon.
@@ -131,7 +162,7 @@ def test_init_no_csp_explicit(self):
 
         importlib.reload(ccflow.enums)
 
-        # csp will only be in sys.modules if it's actually installed
-        if importlib.util.find_spec("csp") is not None:
+        # csp will only be in sys.modules if it's actually installed and older than 0.19
+        if importlib.util.find_spec("csp") is not None and ccflow.enums._csp_version() < (0, 19):
             self.assertTrue("csp" in sys.modules)
         os.environ.pop("CCFLOW_NO_CSP", None)

ccflow/tests/models/test_narwhals.py

@@ -243,7 +243,7 @@ def test_basic_join(self):
         lf = nw.from_native(pl.LazyFrame({"x": [1, 2, 3]}))
         right_src = FrameSource(data={"x": [1, 2, 3], "y": ["a", "b", "c"]})
         jt = JoinTransform(other=right_src, on="x", how="left")
-        out = lf.pipe(jt).collect().to_native()
+        out = lf.pipe(jt).collect().to_native().sort("x")
         assert out.columns == ["x", "y"]
         assert out["y"].to_list() == ["a", "b", "c"]
 
@@ -255,7 +255,7 @@ def test_json_roundtrip(self):
         jt = JoinTransform(other=right_src, on="x", how="left", suffix="_r")
         jt2 = JoinTransform.model_validate_json(jt.model_dump_json())
         lf = nw.from_native(pl.LazyFrame({"x": [1, 2, 3]}))
-        assert lf.pipe(jt).collect().to_native().equals(lf.pipe(jt2).collect().to_native())
+        assert lf.pipe(jt).collect().to_native().sort("x").equals(lf.pipe(jt2).collect().to_native().sort("x"))
 
     def test_rejects_non_narwhals_other(self):
         class BadSrc(CallableModel):
@@ -271,15 +271,15 @@ def test_in_pipeline(self):
         main = FrameSource(data={"x": [1, 2, 3], "v": [10, 20, 30]})
         side = FrameSource(data={"x": [1, 2, 3], "label": ["a", "b", "c"]})
         p = NarwhalsPipelineModel(source=main, transforms=[JoinTransform(other=side, on="x")])
-        out = p(NullContext()).df.collect().to_native()
+        out = p(NullContext()).df.collect().to_native().sort("x")
         assert out.columns == ["x", "v", "label"]
         assert out["label"].to_list() == ["a", "b", "c"]
 
     def test_left_on_right_on(self):
         lf = nw.from_native(pl.LazyFrame({"a": [1, 2, 3]}))
         right_src = FrameSource(data={"b": [1, 2, 3], "y": ["a", "b", "c"]})
         jt = JoinTransform(other=right_src, left_on="a", right_on="b", how="inner")
-        out = lf.pipe(jt).collect().to_native()
+        out = lf.pipe(jt).collect().to_native().sort("a")
         assert out["y"].to_list() == ["a", "b", "c"]
 
     def test_cross_join(self):
@@ -318,7 +318,7 @@ def __call__(self, context: TableCtx) -> NarwhalsFrameResult:
         src = TableSource(tables={"r": {"x": [1, 2], "y": ["a", "b"]}})
         lf = nw.from_native(pl.LazyFrame({"x": [1, 2]}))
         jt = JoinTransform(other=src, other_context=TableCtx(table="r"), on="x")
-        out = lf.pipe(jt).collect().to_native()
+        out = lf.pipe(jt).collect().to_native().sort("x")
         assert out["y"].to_list() == ["a", "b"]
 
     def test_rejects_other_context_type_mismatch(self):
@@ -346,7 +346,7 @@ def __call__(self, df: nw.LazyFrame) -> nw.LazyFrame:
                 return df.select(nw.col("x"), (nw.col("v") * 100).alias("hund"))
 
         jb = JoinBackTransform(inner=Project(), on="x")
-        out = lf.pipe(jb).collect().to_native()
+        out = lf.pipe(jb).collect().to_native().sort("x")
         assert out.columns == ["x", "v", "hund"]
         assert out["hund"].to_list() == [1000, 2000, 3000]
 

ccflow/tests/test_context.py

@@ -102,7 +102,7 @@ def test_datetime_validation(self):
         dt = datetime(2022, 1, 1, 12, 0, tzinfo=timezone.utc)
         c = DatetimeContext(dt=dt)
         self.assertEqual(DatetimeContext(dt=str(dt)), c)
-        self.assertEqual(DatetimeContext(dt=dt.date()), DatetimeContext(dt=datetime(2022, 1, 1)))  # noqa: DTZ001
+        self.assertEqual(DatetimeContext(dt=dt.date()), DatetimeContext(dt=datetime(2022, 1, 1)))
 
         # Test validation from other types
         self.assertEqual(TypeAdapter(DatetimeContext).validate_python({"dt": dt}), c)

ccflow/tests/test_validators.py

@@ -18,16 +18,16 @@ def test_normalize_date(self):
         c1 = datetime.now().astimezone().date() - timedelta(1)
         self.assertEqual(normalize_date("-1d"), c1)
 
-        self.assertEqual(normalize_date(datetime.now()), c)  # noqa: DTZ005
-        self.assertEqual(normalize_date(datetime.now().isoformat()), c)  # noqa: DTZ005
+        self.assertEqual(normalize_date(datetime.now()), c)
+        self.assertEqual(normalize_date(datetime.now().isoformat()), c)
 
         self.assertEqual(normalize_date("foo"), "foo")
         self.assertEqual(normalize_date(None), None)
 
     def test_normalize_datetime(self):
         today = datetime.today()  # noqa: DTZ002
-        now = datetime.now()  # noqa: DTZ005
-        c = datetime(today.year, today.month, today.day)  # noqa: DTZ001
+        now = datetime.now()
+        c = datetime(today.year, today.month, today.day)
 
         self.assertEqual(normalize_datetime(c), c)
         self.assertEqual(normalize_datetime("0d"), c)

ccflow/tests/utils/test_tokenize.py

@@ -636,7 +636,7 @@ def test_primitives(self, value, expected):
         "value,expected",
         [
             (date(2024, 1, 15), ("date", "2024-01-15")),
-            (datetime(2024, 1, 15, 10, 30, 0), ("datetime", "2024-01-15T10:30:00")),  # noqa: DTZ001
+            (datetime(2024, 1, 15, 10, 30, 0), ("datetime", "2024-01-15T10:30:00")),
             (time(10, 30, 0), ("time", "10:30:00")),
             (timedelta(hours=1, minutes=30), ("timedelta", 0, 5400, 0)),
         ],