refactor(python): Add TYPE_CHECKING lints (pola-rs#7070)

py-polars/polars/_html.py

@@ -4,10 +4,11 @@
 import html
 import os
 from textwrap import dedent
-from types import TracebackType
 from typing import TYPE_CHECKING, Iterable
 
 if TYPE_CHECKING:
+    from types import TracebackType
+
     from polars.internals import DataFrame
 
 

py-polars/polars/cfg.py

@@ -1,22 +1,22 @@
 from __future__ import annotations
 
+import contextlib
 import json
 import os
-import sys
-from types import TracebackType
+from typing import TYPE_CHECKING
 
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal
-
-try:
+with contextlib.suppress(ImportError):  # Module not available when building docs
     from polars.polars import set_float_fmt as _set_float_fmt
 
-    _DOCUMENTING = False
-except ImportError:
-    _DOCUMENTING = True
 
+if TYPE_CHECKING:
+    import sys
+    from types import TracebackType
+
+    if sys.version_info >= (3, 8):
+        from typing import Literal
+    else:
+        from typing_extensions import Literal
 
 # note: register all Config-specific environment variable names here; need to constrain
 # which 'POLARS_' environment variables are recognised, as there are other lower-level

py-polars/polars/datatypes.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import ctypes
 import functools
 import re
@@ -27,20 +28,13 @@
 from polars.dependencies import numpy as np
 from polars.dependencies import pyarrow as pa
 
-try:
+with contextlib.suppress(ImportError):  # Module not available when building docs
     from polars.polars import dtype_str_repr
     from polars.polars import get_idx_type as _get_idx_type
 
-    _DOCUMENTING = False
-except ImportError:
-    _DOCUMENTING = True
-
-
 if sys.version_info >= (3, 8):
-    from typing import Literal, get_args
+    from typing import get_args
 else:
-    from typing_extensions import Literal
-
     # pass-through (only impact is that under 3.7 we'll end-up doing
     # standard inference for dataclass fields with an option/union)
     def get_args(tp: Any) -> Any:
@@ -50,10 +44,7 @@ def get_args(tp: Any) -> Any:
 OptionType = type(Optional[type])
 if sys.version_info >= (3, 10):
     from types import NoneType, UnionType
-    from typing import TypeAlias
 else:
-    from typing_extensions import TypeAlias
-
     # infer equivalent class
     NoneType = type(None)
     UnionType = type(Union[int, float])
@@ -62,6 +53,15 @@ def get_args(tp: Any) -> Any:
 if TYPE_CHECKING:
     from polars.internals.type_aliases import TimeUnit
 
+    if sys.version_info >= (3, 8):
+        from typing import Literal
+    else:
+        from typing_extensions import Literal
+
+    if sys.version_info >= (3, 10):
+        from typing import TypeAlias
+    else:
+        from typing_extensions import TypeAlias
 
 # number of rows to scan by default when inferring datatypes
 N_INFER_DEFAULT = 100

py-polars/polars/datatypes_constructor.py

@@ -28,6 +28,7 @@
 )
 from polars.dependencies import numpy as np
 
+# Module not available when building docs
 try:
     from polars.polars import PySeries
 

py-polars/polars/internals/anonymous_scan.py

@@ -6,7 +6,7 @@
 
 import polars as pl
 from polars import internals as pli
-from polars.dependencies import pyarrow as pa
+from polars.dependencies import pyarrow as pa  # noqa: TCH001
 
 
 def _deser_and_exec(  # noqa: D417

py-polars/polars/internals/batched.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import contextlib
 from pathlib import Path
-from typing import Sequence
+from typing import TYPE_CHECKING, Sequence
 
 import polars.internals as pli
 from polars.datatypes import (
@@ -10,20 +11,18 @@
     SchemaDict,
     py_type_to_dtype,
 )
-from polars.internals.type_aliases import CsvEncoding
 from polars.utils import (
     _prepare_row_count_args,
     _process_null_values,
     handle_projection_columns,
     normalise_filepath,
 )
 
-try:
+with contextlib.suppress(ImportError):  # Module not available when building docs
     from polars.polars import PyBatchedCsv
 
-    _DOCUMENTING = False
-except ImportError:
-    _DOCUMENTING = True
+if TYPE_CHECKING:
+    from polars.internals.type_aliases import CsvEncoding
 
 
 class BatchedCsvReader:

py-polars/polars/internals/construction.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 from contextlib import suppress
 from dataclasses import astuple, is_dataclass
 from datetime import date, datetime, time, timedelta
@@ -62,6 +63,12 @@
     threadpool_size,
 )
 
+with contextlib.suppress(ImportError):  # Module not available when building docs
+    from polars.polars import PyDataFrame, PySeries
+
+if TYPE_CHECKING:
+    from polars.internals.type_aliases import Orientation
+
 if version_info >= (3, 10):
 
     def dataclass_type_hints(obj: type) -> dict[str, Any]:
@@ -73,18 +80,6 @@ def dataclass_type_hints(obj: type) -> dict[str, Any]:
         return obj.__annotations__
 
 
-try:
-    from polars.polars import PyDataFrame, PySeries
-
-    _DOCUMENTING = False
-except ImportError:
-    _DOCUMENTING = True
-
-
-if TYPE_CHECKING:
-    from polars.internals.type_aliases import Orientation
-
-
 def is_namedtuple(value: Any, annotated: bool = False) -> bool:
     """Infer whether value is a NamedTuple."""
     if all(hasattr(value, attr) for attr in ("_fields", "_field_defaults", "_replace")):

py-polars/polars/internals/dataframe/frame.py

@@ -1,13 +1,12 @@
 """Module containing logic related to eager DataFrames."""
 from __future__ import annotations
 
+import contextlib
 import math
 import os
 import random
-import sys
 import typing
 from collections.abc import Sized
-from datetime import timedelta
 from io import BytesIO, IOBase, StringIO
 from pathlib import Path
 from typing import (
@@ -90,29 +89,14 @@
     scale_bytes,
 )
 
-try:
+with contextlib.suppress(ImportError):  # Module not available when building docs
     from polars.polars import PyDataFrame
 
-    _DOCUMENTING = False
-except ImportError:
-    _DOCUMENTING = True
-
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal
-
-if sys.version_info >= (3, 10):
-    from typing import Concatenate, ParamSpec, TypeAlias
-else:
-    from typing_extensions import Concatenate, ParamSpec, TypeAlias
-
-if sys.version_info >= (3, 11):
-    from typing import Self
-else:
-    from typing_extensions import Self
 
 if TYPE_CHECKING:
+    import sys
+    from datetime import timedelta
+
     from pyarrow.interchange.dataframe import _PyArrowDataFrame
 
     from polars.internals.type_aliases import (
@@ -137,6 +121,21 @@
         UnstackDirection,
     )
 
+    if sys.version_info >= (3, 8):
+        from typing import Literal
+    else:
+        from typing_extensions import Literal
+
+    if sys.version_info >= (3, 10):
+        from typing import Concatenate, ParamSpec, TypeAlias
+    else:
+        from typing_extensions import Concatenate, ParamSpec, TypeAlias
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
     # these aliases are used to annotate DataFrame.__getitem__()
     # MultiRowSelector indexes into the vertical axis and
     # MultiColSelector indexes into the horizontal axis
@@ -6929,8 +6928,8 @@ def row(
         one row is returned; more than one row raises ``TooManyRowsReturned``, and
         zero rows will raise ``NoRowsReturned`` (both inherit from ``RowsException``).
 
-        Warning
-        -------
+        Warnings
+        --------
         You should NEVER use this method to iterate over a DataFrame; if you absolutely
         require row-iteration you should strongly prefer ``iter_rows()`` instead.
 

py-polars/polars/internals/dataframe/groupby.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from datetime import timedelta
 from typing import (
     TYPE_CHECKING,
     Callable,
@@ -14,6 +13,8 @@
 from polars.utils import _timedelta_to_pl_duration, deprecated_alias, redirect
 
 if TYPE_CHECKING:
+    from datetime import timedelta
+
     from polars.internals.type_aliases import (
         ClosedInterval,
         IntoExpr,

py-polars/polars/internals/expr/expr.py

@@ -3,7 +3,6 @@
 import math
 import os
 import random
-import sys
 import warnings
 from datetime import date, datetime, time, timedelta
 from typing import TYPE_CHECKING, Any, Callable, Iterable, NoReturn, Sequence, cast
@@ -26,40 +25,36 @@
 from polars.internals.expr.meta import ExprMetaNameSpace
 from polars.internals.expr.string import ExprStringNameSpace
 from polars.internals.expr.struct import ExprStructNameSpace
-from polars.internals.type_aliases import PythonLiteral
 from polars.utils import (
     _timedelta_to_pl_duration,
     deprecate_nonkeyword_arguments,
     deprecated_alias,
     sphinx_accessor,
 )
 
-try:
-    from polars.polars import PyExpr
-
-    _DOCUMENTING = False
-except ImportError:
-    _DOCUMENTING = True
-
 if TYPE_CHECKING:
+    import sys
+
     from polars.internals.type_aliases import (
         ClosedInterval,
         FillNullStrategy,
         InterpolationMethod,
         IntoExpr,
         NullBehavior,
+        PythonLiteral,
         RankMethod,
         RollingInterpolationMethod,
         SearchSortedSide,
     )
+    from polars.polars import PyExpr
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
 elif os.getenv("BUILDING_SPHINX_DOCS"):
     property = sphinx_accessor
 
-if sys.version_info >= (3, 11):
-    from typing import Self
-else:
-    from typing_extensions import Self
-
 
 def selection_to_pyexpr_list(
     exprs: IntoExpr | Iterable[IntoExpr],

py-polars/polars/internals/expr/list.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
 import copy
-from datetime import date, datetime, time
 from typing import TYPE_CHECKING, Any, Callable
 
 import polars.internals as pli
 from polars.utils import deprecate_nonkeyword_arguments, deprecated_alias
 
 if TYPE_CHECKING:
+    from datetime import date, datetime, time
+
     from polars.internals.type_aliases import NullBehavior, ToStructStrategy