diff --git a/ccflow/enums.py b/ccflow/enums.py index 27653ded..3cc1a1d0 100644 --- a/ccflow/enums.py +++ b/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 diff --git a/ccflow/examples/narwhals_pipelines.ipynb b/ccflow/examples/narwhals_pipelines.ipynb index 7197c320..31b521eb 100644 --- a/ccflow/examples/narwhals_pipelines.ipynb +++ b/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", diff --git a/ccflow/tests/enums/test_enums.py b/ccflow/tests/enums/test_enums.py index f799ece4..1ade6a6d 100644 --- a/ccflow/tests/enums/test_enums.py +++ b/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) diff --git a/ccflow/tests/models/test_narwhals.py b/ccflow/tests/models/test_narwhals.py index d74bc9d7..51d89602 100644 --- a/ccflow/tests/models/test_narwhals.py +++ b/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,7 +271,7 @@ 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"] @@ -279,7 +279,7 @@ 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] diff --git a/ccflow/tests/test_context.py b/ccflow/tests/test_context.py index e1a79c02..c3433fa5 100644 --- a/ccflow/tests/test_context.py +++ b/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) diff --git a/ccflow/tests/test_validators.py b/ccflow/tests/test_validators.py index c4e79e55..6f00cd96 100644 --- a/ccflow/tests/test_validators.py +++ b/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) diff --git a/ccflow/tests/utils/test_tokenize.py b/ccflow/tests/utils/test_tokenize.py index 98343a85..1719c2a6 100644 --- a/ccflow/tests/utils/test_tokenize.py +++ b/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)), ],