py39, py310 lint

pola-rs · Dec 6, 2024 · 8b60706 · 8b60706
1 parent ed030a3
commit 8b60706
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 12 deletions.
diff --git a/py-polars/tests/unit/datatypes/test_enum.py b/py-polars/tests/unit/datatypes/test_enum.py
@@ -20,6 +20,13 @@
 )
 from polars.testing import assert_frame_equal, assert_series_equal
 
+if sys.version_info >= (3, 11):
+    from enum import StrEnum
+
+    PyStrEnum: type[enum.Enum] | None = StrEnum
+else:
+    PyStrEnum = None
+
 
 def test_enum_creation() -> None:
     dtype = pl.Enum(["a", "b"])
@@ -617,14 +624,9 @@ def test_enum_19269() -> None:
 @pytest.mark.parametrize(
     "EnumBase",
     [
-        (enum.Enum,),
-        (enum.StrEnum,),
-        (str, enum.Enum),
-    ]
-    if sys.version_info >= (3, 11)
-    else [
         (enum.Enum,),
         (str, enum.Enum),
+        *([(PyStrEnum,)] if PyStrEnum is not None else []),
     ],
 )
 def test_init_frame_from_enums(EnumBase: tuple[type, ...]) -> None:

diff --git a/py-polars/tests/unit/functions/test_lit.py b/py-polars/tests/unit/functions/test_lit.py
@@ -19,6 +19,14 @@
     from polars._typing import PolarsDataType
 
 
+if sys.version_info >= (3, 11):
+    from enum import StrEnum
+
+    PyStrEnum: type[enum.Enum] | None = StrEnum
+else:
+    PyStrEnum = None
+
+
 @pytest.mark.parametrize(
     "input",
     [
@@ -109,14 +117,9 @@ def test_lit_unsupported_type() -> None:
 @pytest.mark.parametrize(
     "EnumBase",
     [
-        (enum.Enum,),
-        (enum.StrEnum,),
-        (str, enum.Enum),
-    ]
-    if sys.version_info >= (3, 11)
-    else [
         (enum.Enum,),
         (str, enum.Enum),
+        *([(PyStrEnum,)] if PyStrEnum is not None else []),
     ],
 )
 def test_lit_enum_input_16668(EnumBase: tuple[type, ...]) -> None: