fix some test lint

pola-rs · Oct 22, 2024 · ea52dce · ea52dce
1 parent 8832963
commit ea52dce
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 11 deletions.
diff --git a/py-polars/polars/schema.py b/py-polars/polars/schema.py
@@ -2,19 +2,25 @@
 
 from collections import OrderedDict
 from collections.abc import Mapping
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
-from polars.datatypes import DataType, is_polars_dtype
+from polars._typing import PythonDataType
+from polars.datatypes import DataType, DataTypeClass, is_polars_dtype
 from polars.datatypes._parse import parse_into_dtype
 
 if TYPE_CHECKING:
+    import sys
     from collections.abc import Iterable
 
-    from polars._typing import PythonDataType
-    from polars.datatypes import DataTypeClass
+    if sys.version_info >= (3, 10):
+        from typing import TypeAlias
+    else:
+        from typing_extensions import TypeAlias
 
 
 BaseSchema = OrderedDict[str, DataType]
+SchemaInitDataType: TypeAlias = Union[DataType, DataTypeClass, PythonDataType]
+
 
 __all__ = ["Schema"]
 
@@ -65,8 +71,8 @@ class Schema(BaseSchema):
     def __init__(
         self,
         schema: (
-            Mapping[str, DataType | DataTypeClass | PythonDataType]
-            | Iterable[tuple[str, DataType | DataTypeClass | PythonDataType]]
+            Mapping[str, SchemaInitDataType]
+            | Iterable[tuple[str, SchemaInitDataType]]
             | None
         ) = None,
         *,

diff --git a/py-polars/tests/unit/dataframe/test_df.py b/py-polars/tests/unit/dataframe/test_df.py
@@ -760,7 +760,7 @@ def test_to_dummies() -> None:
             "i": [1, 2, 3],
             "category": ["dog", "cat", "cat"],
         },
-        schema={"i": pl.Int32, "category": pl.Categorical},
+        schema={"i": pl.Int32, "category": pl.Categorical("lexical")},
     )
     expected = pl.DataFrame(
         {

diff --git a/py-polars/tests/unit/io/test_spreadsheet.py b/py-polars/tests/unit/io/test_spreadsheet.py
@@ -322,11 +322,11 @@ def test_read_mixed_dtype_columns(
 ) -> None:
     spreadsheet_path = request.getfixturevalue(source)
     schema_overrides = {
-        "Employee ID": pl.Utf8,
-        "Employee Name": pl.Utf8,
-        "Date": pl.Date,
+        "Employee ID": pl.Utf8(),
+        "Employee Name": pl.Utf8(),
+        "Date": pl.Date(),
         "Details": pl.Categorical("lexical"),
-        "Asset ID": pl.Utf8,
+        "Asset ID": pl.Utf8(),
     }
     df = read_spreadsheet(
         spreadsheet_path,