Modern style type hints for the test suite (#3949)

py-polars/tests/conftest.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 
 import pytest

py-polars/tests/io/test_avro.py

@@ -1,4 +1,6 @@
 # flake8: noqa: W191,E101
+from __future__ import annotations
+
 import os
 from typing import List
 
@@ -14,11 +16,11 @@ def example_df() -> pl.DataFrame:
 
 
 @pytest.fixture
-def compressions() -> List[str]:
+def compressions() -> list[str]:
     return ["uncompressed", "snappy", "deflate"]
 
 
-def test_from_to_buffer(example_df: pl.DataFrame, compressions: List[str]) -> None:
+def test_from_to_buffer(example_df: pl.DataFrame, compressions: list[str]) -> None:
     for compression in compressions:
         buf = io.BytesIO()
         example_df.write_avro(buf, compression=compression)  # type: ignore
@@ -28,7 +30,7 @@ def test_from_to_buffer(example_df: pl.DataFrame, compressions: List[str]) -> No
 
 
 def test_from_to_file(
-    io_test_dir: str, example_df: pl.DataFrame, compressions: List[str]
+    io_test_dir: str, example_df: pl.DataFrame, compressions: list[str]
 ) -> None:
     f = os.path.join(io_test_dir, "small.avro")
 

py-polars/tests/io/test_csv.py

@@ -1,4 +1,6 @@
 # flake8: noqa: W191,E101
+from __future__ import annotations
+
 import gzip
 import io
 import os
@@ -166,7 +168,7 @@ def test_partial_column_rename() -> None:
     "col_input, col_out", [([0, 1], ["a", "b"]), ([0, 2], ["a", "c"]), (["b"], ["b"])]
 )
 def test_read_csv_columns_argument(
-    col_input: Union[List[int], List[str]], col_out: List[str]
+    col_input: list[int] | list[str], col_out: list[str]
 ) -> None:
     csv = textwrap.dedent(
         """\
@@ -335,7 +337,7 @@ def test_ignore_parse_dates() -> None:
     ).encode()
 
     headers = ["a", "b", "c"]
-    dtypes: Dict[str, Type[DataType]] = {
+    dtypes: dict[str, type[DataType]] = {
         k: pl.Utf8 for k in headers
     }  # Forces Utf8 type for every column
     df = pl.read_csv(csv, columns=headers, dtypes=dtypes)

py-polars/tests/io/test_ipc.py

@@ -1,4 +1,6 @@
 # flake8: noqa: W191,E101
+from __future__ import annotations
+
 import io
 import os
 from pathlib import Path
@@ -11,11 +13,11 @@
 
 
 @pytest.fixture
-def compressions() -> List[str]:
+def compressions() -> list[str]:
     return ["uncompressed", "lz4", "zstd"]
 
 
-def test_from_to_buffer(df: pl.DataFrame, compressions: List[str]) -> None:
+def test_from_to_buffer(df: pl.DataFrame, compressions: list[str]) -> None:
     for compression in compressions:
         buf = io.BytesIO()
         df.write_ipc(buf, compression=compression)  # type: ignore
@@ -25,7 +27,7 @@ def test_from_to_buffer(df: pl.DataFrame, compressions: List[str]) -> None:
 
 
 def test_from_to_file(
-    io_test_dir: str, df: pl.DataFrame, compressions: List[str]
+    io_test_dir: str, df: pl.DataFrame, compressions: list[str]
 ) -> None:
     f_ipc = os.path.join(io_test_dir, "small.ipc")
 
@@ -74,7 +76,7 @@ def test_compressed_simple() -> None:
         assert df_read.frame_equal(df)
 
 
-def test_ipc_schema(compressions: List[str]) -> None:
+def test_ipc_schema(compressions: list[str]) -> None:
     df = pl.DataFrame({"a": [1, 2], "b": ["a", None], "c": [True, False]})
 
     for compression in compressions:
@@ -86,7 +88,7 @@ def test_ipc_schema(compressions: List[str]) -> None:
 
 
 def test_ipc_schema_from_file(
-    io_test_dir: str, df_no_lists: pl.DataFrame, compressions: List[str]
+    io_test_dir: str, df_no_lists: pl.DataFrame, compressions: list[str]
 ) -> None:
     df = df_no_lists
     f_ipc = os.path.join(io_test_dir, "small.ipc")

py-polars/tests/io/test_json.py

@@ -1,4 +1,6 @@
 # flake8: noqa: W191,E101
+from __future__ import annotations
+
 import io
 import os
 

py-polars/tests/io/test_lazy_csv.py

@@ -1,4 +1,6 @@
 # flake8: noqa: W191,E101
+from __future__ import annotations
+
 import io
 from os import path
 from pathlib import Path

py-polars/tests/io/test_lazy_ipc.py

@@ -1,4 +1,6 @@
 # flake8: noqa: W191,E101
+from __future__ import annotations
+
 import polars as pl
 
 

py-polars/tests/io/test_lazy_parquet.py

@@ -1,4 +1,6 @@
 # flake8: noqa: W191,E101
+from __future__ import annotations
+
 from os import path
 from pathlib import Path
 

py-polars/tests/io/test_other.py

@@ -1,4 +1,6 @@
 # flake8: noqa: W191,E101
+from __future__ import annotations
+
 import copy
 
 import polars as pl

py-polars/tests/io/test_parquet.py

@@ -1,4 +1,6 @@
 # flake8: noqa: W191,E101
+from __future__ import annotations
+
 import io
 import os
 from typing import List
@@ -11,11 +13,11 @@
 
 
 @pytest.fixture
-def compressions() -> List[str]:
+def compressions() -> list[str]:
     return ["uncompressed", "snappy", "gzip", "lzo", "brotli", "lz4", "zstd"]
 
 
-def test_to_from_buffer(df: pl.DataFrame, compressions: List[str]) -> None:
+def test_to_from_buffer(df: pl.DataFrame, compressions: list[str]) -> None:
     for compression in compressions:
         if compression == "lzo":
             # lzo compression is not supported now
@@ -39,7 +41,7 @@ def test_to_from_buffer(df: pl.DataFrame, compressions: List[str]) -> None:
 
 
 def test_to_from_file(
-    io_test_dir: str, df: pl.DataFrame, compressions: List[str]
+    io_test_dir: str, df: pl.DataFrame, compressions: list[str]
 ) -> None:
     f = os.path.join(io_test_dir, "small.parquet")
     for compression in compressions:

py-polars/tests/io/test_pickle.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pickle
 
 import polars as pl

py-polars/tests/io/test_pyarrow_dataset.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 
 import pyarrow.dataset as ds

py-polars/tests/io/test_sql.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 from datetime import date
 

py-polars/tests/run_doc_examples.py

@@ -18,6 +18,8 @@
 facilitate such a change, whilst not immediately having to add IGNORE_RESULT directives everywhere or changing all
 outputs, set `IGNORE_RESULT_ALL=True` below. Do note that this does mean no output is being checked anymore.
 """
+from __future__ import annotations
+
 import doctest
 import importlib
 import sys

py-polars/tests/test_apply.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import typing
 from datetime import date, datetime, timedelta
 from functools import reduce
-from typing import List, Optional
 
 import numpy as np
 
@@ -32,7 +33,7 @@ def test_apply_none() -> None:
     assert out_df["a"].to_list() == (df["a"] * df["b"]).to_list()
 
     # check if we can return None
-    def func(s: List) -> Optional[int]:
+    def func(s: list) -> int | None:
         if s[0][0] == 190:
             return None
         else:

py-polars/tests/test_categorical.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import io
 
 import polars as pl

py-polars/tests/test_cfg.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 from typing import Generator
 

py-polars/tests/test_datatypes.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import inspect
 
 import polars as pl
@@ -6,11 +8,11 @@
 
 def test_dtype_init_equivalence() -> None:
     # check "DataType.__new__" behaviour for all datatypes
-    all_datatypes = set(
+    all_datatypes = {
         dtype
         for dtype in (getattr(datatypes, attr) for attr in dir(datatypes))
         if inspect.isclass(dtype) and issubclass(dtype, datatypes.DataType)
-    )
+    }
     for dtype in all_datatypes:
         assert dtype == dtype()
 

py-polars/tests/test_datelike.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import io
 import typing
 from datetime import date, datetime, time, timedelta
@@ -985,7 +987,7 @@ def test_datetime_units() -> None:
     names = set(df.columns)
 
     for unit in DTYPE_TEMPORAL_UNITS:
-        subset = names - set([unit])
+        subset = names - {unit}
 
         assert (
             len(set(df.select([pl.all().exclude(pl.Datetime(unit))]).columns) - subset)

py-polars/tests/test_df.py

@@ -1,4 +1,6 @@
 # flake8: noqa: W191,E101
+from __future__ import annotations
+
 import sys
 import typing
 from builtins import range
@@ -2172,7 +2174,7 @@ class SubClassedDataFrame(pl.DataFrame):
     # and connect these classes together
     class MyLazyFrame(pl.LazyFrame):
         @property
-        def _dataframe_class(cls) -> "Type[MyDataFrame]":
+        def _dataframe_class(cls) -> type[MyDataFrame]:
             return MyDataFrame
 
     class MyDataFrame(pl.DataFrame):

py-polars/tests/test_errors.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import io
 import typing
 

py-polars/tests/test_explode.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pyarrow as pa
 
 import polars as pl

py-polars/tests/test_exprs.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import polars as pl
 from polars.testing import assert_series_equal, verify_series_and_expr_api
 

py-polars/tests/test_functions.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import polars as pl
 
 

py-polars/tests/test_groupby.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from datetime import datetime
 
 import polars as pl

py-polars/tests/test_interop.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import typing
 from datetime import datetime
-from typing import Dict, Sequence, Type, Union
+from typing import Sequence
 from unittest.mock import patch
 
 import numpy as np
@@ -53,13 +55,13 @@ def test_to_numpy() -> None:
     def test_series_to_numpy(
         name: str,
         values: list,
-        pl_dtype: Type[pl.DataType],
-        np_dtype: Union[
-            Type[np.signedinteger],
-            Type[np.unsignedinteger],
-            Type[np.floating],
-            Type[np.object_],
-        ],
+        pl_dtype: type[pl.DataType],
+        np_dtype: (
+            type[np.signedinteger]
+            | type[np.unsignedinteger]
+            | type[np.floating]
+            | type[np.object_]
+        ),
     ) -> None:
         pl_series_to_numpy_array = np.array(pl.Series(name, values, pl_dtype))
         numpy_array = np.array(values, dtype=np_dtype)
@@ -216,7 +218,7 @@ def test_from_dict() -> None:
 
 
 def test_from_dict_struct() -> None:
-    data: Dict[str, Union[Dict, Sequence]] = {
+    data: dict[str, dict | Sequence] = {
         "a": {"b": [1, 3], "c": [2, 4]},
         "d": [5, 6],
     }

py-polars/tests/test_joins.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from datetime import datetime
 
 import numpy as np

py-polars/tests/test_lazy.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import numpy as np
 import pytest
 from _pytest.capture import CaptureFixture

py-polars/tests/test_lists.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from datetime import date, datetime, time
 
 import pandas as pd
@@ -168,7 +170,7 @@ def test_list_append() -> None:
     out = df.select([pl.col("a").arr.concat([1, 4])])
     assert out["a"][0].to_list() == [1, 2, 1, 4]
 
-    out_s = df["a"].arr.concat(([4, 1]))
+    out_s = df["a"].arr.concat([4, 1])
     assert out_s[0].to_list() == [1, 2, 4, 1]
 
 

py-polars/tests/test_queries.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from datetime import datetime
 
 import numpy as np
@@ -25,8 +27,8 @@ def test_type_coercion_when_then_otherwise_2806() -> None:
         pl.DataFrame({"names": ["foo", "spam", "spam"], "nrs": [1, 2, 3]})
         .select(
             [
-                pl.when((pl.col("names") == "spam"))
-                .then((pl.col("nrs") * 2))
+                pl.when(pl.col("names") == "spam")
+                .then(pl.col("nrs") * 2)
                 .otherwise(pl.lit("other"))
                 .alias("new_col"),
             ]