Enable, fix ruff TC checks

doc/conf.py

@@ -46,8 +46,13 @@
 
 sys._BUILDING_SPHINX_DOCS = True
 
+
 nitpick_ignore_regex = [
     ["py:class", r"numpy.(u?)int[\d]+"],
     ["py:class", r"numpy.bool_"],
     ["py:class", r"typing_extensions(.+)"],
+    ["py:class", r"P\.args"],
+    ["py:class", r"P\.kwargs"],
+    ["py:class", r"lp\.LoopKernel"],
+    ["py:class", r"_dtype_any"],
 ]

pyproject.toml

@@ -62,6 +62,7 @@ extend-select = [
     "NPY",   # numpy
     "RUF",
     "UP",
+    "TC",
 ]
 extend-ignore = [
   "E226",

pytato/analysis/__init__.py

@@ -26,7 +26,6 @@
 THE SOFTWARE.
 """
 
-from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any
 
 from loopy.tools import LoopyKeyBuilder
@@ -47,12 +46,14 @@
     Stack,
 )
 from pytato.function import Call, FunctionDefinition, NamedCallResult
-from pytato.loopy import LoopyCall
 from pytato.transform import ArrayOrNames, CachedWalkMapper, Mapper
 
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
+
     from pytato.distributed.nodes import DistributedRecv, DistributedSendRefHolder
+    from pytato.loopy import LoopyCall
 
 __doc__ = """
 .. currentmodule:: pytato.analysis

pytato/array.py

@@ -405,7 +405,7 @@ def _dataclass_setstate(self, state):
 
     # {{{ assign mapper_method
 
-    mm_cls = cast(type[_HasMapperMethod], cls)
+    mm_cls = cast("type[_HasMapperMethod]", cls)
 
     snake_clsname = _CAMEL_TO_SNAKE_RE.sub("_", mm_cls.__name__).lower()
     default_mapper_method_name = f"map_{snake_clsname}"
@@ -804,7 +804,7 @@ def conj(self) -> ArrayOrScalar:
 
     def __abs__(self) -> Array:
         import pytato as pt
-        return cast(Array, pt.abs(self))
+        return cast("Array", pt.abs(self))
 
     def __pos__(self) -> Array:
         return self
@@ -1755,7 +1755,7 @@ def shape(self) -> ShapeType:
                        for i_basic_idx in i_basic_indices)
 
         adv_idx_shape = get_shape_after_broadcasting([
-            cast(Array | Integer, not_none(self.indices[i_idx]))
+            cast("Array | Integer", not_none(self.indices[i_idx]))
             for i_idx in i_adv_indices])
 
         # type-ignored because mypy cannot figure out basic-indices only refer
@@ -1803,7 +1803,7 @@ def shape(self) -> ShapeType:
                    for i_basic_idx in i_basic_indices)
 
         adv_idx_shape = get_shape_after_broadcasting([
-            cast(Array | Integer, not_none(self.indices[i_idx]))
+            cast("Array | Integer", not_none(self.indices[i_idx]))
             for i_idx in i_adv_indices])
 
         # type-ignored because mypy cannot figure out basic-indices only refer slices
@@ -2037,7 +2037,7 @@ def matmul(x1: Array, x2: Array) -> Array:
     if x1.ndim == x2.ndim == 1:
         return pt.sum(x1 * x2)
     elif x1.ndim == 1:
-        return cast(Array, pt.dot(x1, x2))
+        return cast("Array", pt.dot(x1, x2))
     elif x2.ndim == 1:
         x1_indices = index_names[:x1.ndim]
         return pt.einsum(f"{x1_indices}, {x1_indices[-1]} -> {x1_indices[:-1]}",
@@ -2370,7 +2370,7 @@ def full(shape: ConvertibleToShape, fill_value: Scalar | prim.NaN,
     else:
         fill_value = conv_dtype.type(fill_value)
 
-    return IndexLambda(expr=cast(ArithmeticExpression, fill_value),
+    return IndexLambda(expr=cast("ArithmeticExpression", fill_value),
                        shape=shape, dtype=conv_dtype,
                        bindings=immutabledict(),
                        tags=_get_default_tags(),

pytato/cmath.py

@@ -57,14 +57,13 @@
 # }}}
 
 
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 from immutabledict import immutabledict
 
 import pymbolic.primitives as prim
 from pymbolic import Scalar, var
-from pymbolic.typing import Expression
 
 from pytato.array import (
     Array,
@@ -78,6 +77,10 @@
 from pytato.scalar_expr import SCALAR_CLASSES
 
 
+if TYPE_CHECKING:
+    from pymbolic.typing import Expression
+
+
 def _apply_elem_wise_func(inputs: tuple[ArrayOrScalar, ...],
                           func_name: str,
                           ret_dtype: _dtype_any | None = None,
@@ -88,7 +91,7 @@ def _apply_elem_wise_func(inputs: tuple[ArrayOrScalar, ...],
             np_func_name = func_name
 
         np_func = getattr(np, np_func_name)
-        return cast(ArrayOrScalar, np_func(*inputs))
+        return cast("ArrayOrScalar", np_func(*inputs))
 
     if not inputs:
         raise ValueError("at least one argument must be present")
@@ -233,7 +236,7 @@ def imag(x: ArrayOrScalar) -> ArrayOrScalar:
         result_dtype = np.empty(0, dtype=x_dtype).real.dtype
     else:
         if np.isscalar(x):
-            return cast(Scalar, x_dtype.type(0))
+            return cast("Scalar", x_dtype.type(0))
         else:
             assert isinstance(x, Array)
             import pytato as pt

pytato/codegen.py

@@ -38,8 +38,7 @@
 """
 
 import dataclasses
-from collections.abc import Mapping
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from immutabledict import immutabledict
 
@@ -57,10 +56,8 @@
     SizeParam,
     make_dict_of_named_arrays,
 )
-from pytato.function import NamedCallResult
 from pytato.loopy import LoopyCall
 from pytato.scalar_expr import IntegralScalarExpression, is_integral_scalar_expression
-from pytato.target import Target
 from pytato.transform import (
     ArrayOrNames,
     CachedWalkMapper,
@@ -70,6 +67,13 @@
 from pytato.transform.lower_to_index_lambda import ToIndexLambdaMixin
 
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from pytato.function import NamedCallResult
+    from pytato.target import Target
+
+
 SymbolicIndex: TypeAlias = tuple[IntegralScalarExpression, ...]
 
 

pytato/distributed/execute.py

@@ -34,28 +34,30 @@
 """
 
 import logging
-from collections.abc import Hashable, Mapping
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 from pytato.array import make_dict_of_named_arrays
-from pytato.distributed.nodes import DistributedRecv, DistributedSend
-from pytato.distributed.partition import (
-    DistributedGraphPart,
-    DistributedGraphPartition,
-    PartId,
-)
 from pytato.scalar_expr import INT_CLASSES
-from pytato.target import BoundProgram
 
 
 logger = logging.getLogger(__name__)
 
 
 if TYPE_CHECKING:
+    from collections.abc import Hashable, Mapping
+
     import mpi4py.MPI
 
+    from pytato.distributed.nodes import DistributedRecv, DistributedSend
+    from pytato.distributed.partition import (
+        DistributedGraphPart,
+        DistributedGraphPartition,
+        PartId,
+    )
+    from pytato.target import BoundProgram
+
 
 # {{{ generate_code_for_partition
 
@@ -134,7 +136,7 @@ def execute_distributed_partition(
     context: dict[str, Any] = input_args.copy()
 
     pids_to_execute = set(partition.parts)
-    pids_executed = set()
+    pids_executed: set[PartId] = set()
     recv_names_completed = set()
     send_requests = []
 

pytato/distributed/partition.py

@@ -88,14 +88,15 @@
     DistributedSend,
     DistributedSendRefHolder,
 )
-from pytato.function import FunctionDefinition, NamedCallResult
 from pytato.scalar_expr import SCALAR_CLASSES
 from pytato.transform import ArrayOrNames, CachedWalkMapper, CombineMapper, CopyMapper
 
 
 if TYPE_CHECKING:
     import mpi4py.MPI
 
+    from pytato.function import FunctionDefinition, NamedCallResult
+
 
 @dataclasses.dataclass(frozen=True)
 class CommunicationOpIdentifier:
@@ -350,7 +351,7 @@ def rec(self, expr: ArrayOrNames) -> ArrayOrNames:
             if name is not None:
                 return self._get_placeholder_for(name, expr)
 
-        return cast(ArrayOrNames, super().rec(expr))
+        return cast("ArrayOrNames", super().rec(expr))
 
     def _get_placeholder_for(self, name: str, expr: Array) -> Placeholder:
         placeholder = self.partition_input_name_to_placeholder.get(name)
@@ -820,7 +821,7 @@ def find_distributed_partition(
             raise comm_batches_or_exc
 
         comm_batches = cast(
-                Sequence[Set[CommunicationOpIdentifier]],
+                "Sequence[Set[CommunicationOpIdentifier]]",
                 comm_batches_or_exc)
 
     # }}}
@@ -921,7 +922,7 @@ def find_distributed_partition(
                 ary: max(
                         (comm_id_to_part_id[
                             _recv_to_comm_id(local_rank,
-                                            cast(DistributedRecv, recvd_ary))]
+                                            cast("DistributedRecv", recvd_ary))]
                         for recvd_ary in recvd_array_dep_mapper(ary)),
                         default=-1)
                 for ary in mso_arrays

pytato/distributed/verify.py

@@ -38,11 +38,8 @@
 
 import dataclasses
 import logging
-from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-import numpy as np
-
 from pymbolic.mapper.optimize import optimize_mapper
 
 from pytato.array import (
@@ -51,7 +48,6 @@
     ShapeType,
     make_dict_of_named_arrays,
 )
-from pytato.distributed.nodes import CommTagType, DistributedRecv
 from pytato.distributed.partition import (
     CommunicationOpIdentifier,
     DistributedGraphPartition,
@@ -64,7 +60,12 @@
 
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     import mpi4py.MPI
+    import numpy as np
+
+    from pytato.distributed.nodes import CommTagType, DistributedRecv
 
 
 # {{{ data structures

pytato/equality.py

@@ -25,7 +25,6 @@
 THE SOFTWARE.
 """
 
-from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
 from pytools import memoize_method
@@ -50,11 +49,13 @@
     SizeParam,
     Stack,
 )
-from pytato.function import Call, FunctionDefinition, NamedCallResult
 
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from pytato.distributed.nodes import DistributedRecv, DistributedSendRefHolder
+    from pytato.function import Call, FunctionDefinition, NamedCallResult
     from pytato.loopy import LoopyCall, LoopyCallResult
 
 __doc__ = """

pytato/loopy.py

@@ -27,9 +27,9 @@
 
 
 import dataclasses
-from collections.abc import Iterable, Iterator, Mapping, Sequence
 from numbers import Number
 from typing import (
+    TYPE_CHECKING,
     Any,
 )
 
@@ -58,6 +58,10 @@
 )
 
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator, Mapping, Sequence
+
+
 __doc__ = r"""
 .. currentmodule:: pytato.loopy
 

pytato/pad.py

@@ -1,26 +1,37 @@
 """
 .. autofunction:: pad
+
+Cross-references
+----------------
+
+.. class:: Integer
+
+    See :mod:`pymbolic.typing`.
 """
 from __future__ import annotations
 
 
 __copyright__ = "Copyright (C) 2023 Kaushik Kulkarni"
 
 import collections.abc as abc
-from collections.abc import Sequence
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 import pymbolic.primitives as prim
 from pymbolic import Scalar
-from pymbolic.typing import Integer
 from pytools import UniqueNameGenerator
 
 from pytato.array import Array, IndexLambda
 from pytato.scalar_expr import INT_CLASSES
 
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from pymbolic.typing import Integer
+
+
 def _get_constant_padded_idx_lambda(
     array: Array,
     pad_widths: Sequence[tuple[Integer, Integer]],