Enable TC ruff rules, fix

inducer · inducer · commit dee37a82bda2 · 2024-12-09T21:04:33.000-06:00
diff --git a/pymbolic/algorithm.py b/pymbolic/algorithm.py
@@ -43,7 +43,11 @@
 from pytools import MovedFunctionDeprecationWrapper, memoize
 
 
-if TYPE_CHECKING or getattr(sys, "_BUILDING_SPHINX_DOCS", None):
+if TYPE_CHECKING:
+    import numpy as np
+
+
+if getattr(sys, "_BUILDING_SPHINX_DOCS", None):
     import numpy as np
 
 
diff --git a/pymbolic/geometric_algebra/__init__.py b/pymbolic/geometric_algebra/__init__.py
@@ -567,7 +567,7 @@ def __init__(
                     f"are supported for 'data': shape {data.shape}")
 
             dimensions, = data.shape
-            data_dict = {(i,): cast(CoeffT, xi) for i, xi in enumerate(data)}
+            data_dict = {(i,): cast("CoeffT", xi) for i, xi in enumerate(data)}
 
             if space is None:
                 space = get_euclidean_space(dimensions)
@@ -579,7 +579,7 @@ def __init__(
         elif isinstance(data, Mapping):
             data_dict = data
         else:
-            data_dict = {0: cast(CoeffT, data)}
+            data_dict = {0: cast("CoeffT", data)}
 
         if space is None:
             raise ValueError("No 'space' provided")
@@ -595,16 +595,16 @@ def __init__(
                 assert isinstance(basis_indices, tuple)
 
                 bits, sign = space.bits_and_sign(basis_indices)
-                new_coeff = cast(CoeffT,
-                    new_data.setdefault(bits, cast(CoeffT, 0))  # type: ignore[operator]
+                new_coeff = cast("CoeffT",
+                    new_data.setdefault(bits, cast("CoeffT", 0))  # type: ignore[operator]
                     + sign*coeff)
 
                 if is_zero(new_coeff):
                     del new_data[bits]
                 else:
                     new_data[bits] = new_coeff
         else:
-            new_data = cast(dict[int, CoeffT], data_dict)
+            new_data = cast("dict[int, CoeffT]", data_dict)
 
         # }}}
 
@@ -691,8 +691,8 @@ def __add__(self, other) -> MultiVector:
         from pymbolic.primitives import is_zero
         new_data = {}
         for bits in all_bits:
-            new_coeff = (self.data.get(bits, cast(CoeffT, 0))
-                + other.data.get(bits, cast(CoeffT, 0)))
+            new_coeff = (self.data.get(bits, cast("CoeffT", 0))
+                + other.data.get(bits, cast("CoeffT", 0)))
 
             if not is_zero(new_coeff):
                 new_data[bits] = new_coeff
@@ -741,7 +741,7 @@ def _generic_product(self,
                     coeff = (weight
                             * canonical_reordering_sign(sbits, obits)
                             * scoeff * ocoeff)
-                    new_coeff = new_data.setdefault(new_bits, cast(CoeffT, 0)) + coeff
+                    new_coeff = new_data.setdefault(new_bits, cast("CoeffT", 0)) + coeff
                     if is_zero(new_coeff):
                         del new_data[new_bits]
                     else:
@@ -1134,7 +1134,7 @@ def componentwise(f: Callable[[CoeffT], CoeffT], expr: T) -> T:
     """
 
     if isinstance(expr, MultiVector):
-        return cast(T, expr.map(f))
+        return cast("T", expr.map(f))
 
     from pytools.obj_array import obj_array_vectorize
     return obj_array_vectorize(f, expr)
diff --git a/pymbolic/geometric_algebra/mapper.py b/pymbolic/geometric_algebra/mapper.py
@@ -25,8 +25,7 @@
 
 # This is experimental, undocumented, and could go away any second.
 # Consider yourself warned.
-from collections.abc import Set
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
 import pymbolic.geometric_algebra.primitives as prim
 from pymbolic.geometric_algebra import MultiVector
@@ -49,7 +48,12 @@
     PREC_NONE,
     StringifyMapper as StringifyMapperBase,
 )
-from pymbolic.primitives import ExpressionNode
+
+
+if TYPE_CHECKING:
+    from collections.abc import Set
+
+    from pymbolic.primitives import ExpressionNode
 
 
 class IdentityMapper(IdentityMapperBase[P]):
diff --git a/pymbolic/geometric_algebra/primitives.py b/pymbolic/geometric_algebra/primitives.py
@@ -26,11 +26,15 @@
 # This is experimental, undocumented, and could go away any second.
 # Consider yourself warned.
 
-from collections.abc import Hashable
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
 from pymbolic.primitives import ExpressionNode, Variable, expr_dataclass
-from pymbolic.typing import Expression
+
+
+if TYPE_CHECKING:
+    from collections.abc import Hashable
+
+    from pymbolic.typing import Expression
 
 
 class MultiVectorVariable(Variable):
diff --git a/pymbolic/interop/ast.py b/pymbolic/interop/ast.py
@@ -27,11 +27,14 @@
 """
 
 import ast
-from typing import Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import pymbolic.primitives as p
 from pymbolic.mapper import CachedMapper
-from pymbolic.typing import Expression
+
+
+if TYPE_CHECKING:
+    from pymbolic.typing import Expression
 
 
 __doc__ = r'''
diff --git a/pymbolic/interop/matchpy/mapper.py b/pymbolic/interop/matchpy/mapper.py
@@ -1,9 +1,12 @@
 from __future__ import annotations
 
-from collections.abc import Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-from pymbolic.interop.matchpy import PymbolicOp
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from pymbolic.interop.matchpy import PymbolicOp
 
 
 class Mapper:
diff --git a/pymbolic/interop/matchpy/tofrom.py b/pymbolic/interop/matchpy/tofrom.py
@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import multiset
 import numpy as np
@@ -12,7 +11,12 @@
 import pymbolic.primitives as p
 from pymbolic.interop.matchpy.mapper import Mapper as BaseMatchPyMapper
 from pymbolic.mapper import Mapper as BasePymMapper
-from pymbolic.typing import Scalar as PbScalar
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from pymbolic.typing import Scalar as PbScalar
 
 
 # {{{ to matchpy
diff --git a/pymbolic/mapper/__init__.py b/pymbolic/mapper/__init__.py
@@ -429,7 +429,7 @@ def __call__(self,
         method_name = getattr(expr, "mapper_method", None)
         if method_name is not None:
             method = cast(
-                Callable[Concatenate[Expression, P], ResultT] | None,
+                "Callable[Concatenate[Expression, P], ResultT] | None",
                 getattr(self, method_name, None)
                 )
             if method is not None:
@@ -973,7 +973,7 @@ def map_multivector(self,
                 *args: P.args, **kwargs: P.kwargs
             ) -> Expression:
         # True fact: MultiVectors aren't expressions
-        return expr.map(lambda ch: cast(ArithmeticExpression,
+        return expr.map(lambda ch: cast("ArithmeticExpression",
                                         self.rec(ch, *args, **kwargs)))  # type: ignore[return-value]
 
     def map_common_subexpression(self,
@@ -1012,7 +1012,7 @@ def map_derivative(self,
     def map_slice(self,
                 expr: p.Slice,
                 *args: P.args, **kwargs: P.kwargs) -> Expression:
-        children: p.SliceChildrenT = cast(p.SliceChildrenT, tuple([
+        children: p.SliceChildrenT = cast("p.SliceChildrenT", tuple([
             None if child is None else self.rec(child, *args, **kwargs)
             for child in expr.children
             ]))
diff --git a/pymbolic/mapper/coefficient.py b/pymbolic/mapper/coefficient.py
@@ -68,7 +68,7 @@ def map_product(self, expr: p.Product) -> CoeffsT:
         for i, child_coeffs in enumerate(children_coeffs):
             if i != idx_of_child_with_vars:
                 assert len(child_coeffs) == 1
-                other_coeffs *= cast(ArithmeticExpression, child_coeffs[1])
+                other_coeffs *= cast("ArithmeticExpression", child_coeffs[1])
 
         if idx_of_child_with_vars is None:
             return {1: other_coeffs}
diff --git a/pymbolic/mapper/collector.py b/pymbolic/mapper/collector.py
@@ -26,14 +26,18 @@
 THE SOFTWARE.
 """
 
-from collections.abc import Sequence, Set
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import pymbolic
 import pymbolic.primitives as p
 from pymbolic.mapper import IdentityMapper
-from pymbolic.mapper.dependency import DependenciesT
-from pymbolic.typing import ArithmeticExpression, Expression
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence, Set
+
+    from pymbolic.mapper.dependency import DependenciesT
+    from pymbolic.typing import ArithmeticExpression, Expression
 
 
 class TermCollector(IdentityMapper[[]]):
@@ -110,7 +114,7 @@ def exponent(term: Expression) -> ArithmeticExpression:
 
         base_exp_set = frozenset(
                 (base, exp) for base, exp in cleaned_base2exp.items())
-        return base_exp_set, cast(ArithmeticExpression,
+        return base_exp_set, cast("ArithmeticExpression",
                 self.rec(pymbolic.flattened_product(coefficients)))
 
     def map_sum(self, expr: p.Sum) -> Expression:
diff --git a/pymbolic/mapper/constant_folder.py b/pymbolic/mapper/constant_folder.py
@@ -27,7 +27,8 @@
 THE SOFTWARE.
 """
 
-from collections.abc import Callable
+
+from typing import TYPE_CHECKING
 
 from pymbolic.mapper import (
     CSECachingMapperMixin,
@@ -38,6 +39,10 @@
 from pymbolic.typing import ArithmeticExpression, Expression
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 class ConstantFoldingMapperBase(Mapper[Expression, []]):
     def is_constant(self, expr):
         from pymbolic.mapper.dependency import DependencyMapper
diff --git a/pymbolic/mapper/distributor.py b/pymbolic/mapper/distributor.py
@@ -27,14 +27,17 @@
 THE SOFTWARE.
 """
 
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import pymbolic
 import pymbolic.primitives as p
 from pymbolic.mapper import IdentityMapper
 from pymbolic.mapper.collector import TermCollector
 from pymbolic.mapper.constant_folder import CommutativeConstantFoldingMapper
-from pymbolic.typing import ArithmeticExpression, Expression
+
+
+if TYPE_CHECKING:
+    from pymbolic.typing import ArithmeticExpression, Expression
 
 
 class DistributeMapper(IdentityMapper[[]]):
@@ -118,7 +121,7 @@ def map_power(self, expr: p.Power) -> Expression:
         newbase = self.rec(expr.base)
         if isinstance(newbase, p.Product):
             return self.rec(pymbolic.flattened_product([
-                cast(ArithmeticExpression, child)**expr.exponent
+                cast("ArithmeticExpression", child)**expr.exponent
                     for child in newbase.children
                 ]))
 
diff --git a/pymbolic/mapper/evaluator.py b/pymbolic/mapper/evaluator.py
@@ -34,19 +34,20 @@
 """
 
 import operator as op
-from collections.abc import Mapping
 from functools import reduce
 from typing import TYPE_CHECKING, cast
 
-import pymbolic.primitives as p
 from pymbolic.mapper import CachedMapper, CSECachingMapperMixin, Mapper, ResultT
-from pymbolic.typing import Expression
 
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
+
     import numpy as np
 
+    import pymbolic.primitives as p
     from pymbolic.geometric_algebra import MultiVector
+    from pymbolic.typing import Expression
 
 
 class UnknownVariableError(Exception):
@@ -82,7 +83,7 @@ def __init__(self, context: Mapping[str, ResultT] | None = None) -> None:
         self.context = context
 
     def map_constant(self, expr: object) -> ResultT:
-        return cast(ResultT, expr)
+        return cast("ResultT", expr)
 
     def map_variable(self, expr: p.Variable) -> ResultT:
         try:
diff --git a/pymbolic/mapper/flattener.py b/pymbolic/mapper/flattener.py
@@ -31,15 +31,18 @@
 THE SOFTWARE.
 """
 
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import pymbolic.primitives as p
 from pymbolic.mapper import IdentityMapper
-from pymbolic.typing import (
-    ArithmeticExpression,
-    ArithmeticOrExpressionT,
-    Expression,
-)
+
+
+if TYPE_CHECKING:
+    from pymbolic.typing import (
+        ArithmeticExpression,
+        ArithmeticOrExpressionT,
+        Expression,
+    )
 
 
 class FlattenMapper(IdentityMapper[[]]):
@@ -68,13 +71,13 @@ def is_expr_integer_valued(self, expr: Expression) -> bool:
     def map_sum(self, expr: p.Sum) -> Expression:
         from pymbolic.primitives import flattened_sum
         return flattened_sum([
-                             cast(ArithmeticExpression, self.rec(ch))
+                             cast("ArithmeticExpression", self.rec(ch))
                              for ch in expr.children])
 
     def map_product(self, expr: p.Product) -> Expression:
         from pymbolic.primitives import flattened_product
         return flattened_product([
-                                 cast(ArithmeticExpression, self.rec(ch))
+                                 cast("ArithmeticExpression", self.rec(ch))
                                  for ch in expr.children])
 
     def map_quotient(self, expr: p.Quotient) -> Expression:
@@ -123,4 +126,4 @@ def map_power(self, expr: p.Power) -> Expression:
 
 
 def flatten(expr: ArithmeticOrExpressionT) -> ArithmeticOrExpressionT:
-    return cast(ArithmeticOrExpressionT, FlattenMapper()(expr))
+    return cast("ArithmeticOrExpressionT", FlattenMapper()(expr))
diff --git a/pymbolic/mapper/optimize.py b/pymbolic/mapper/optimize.py
diff --git a/pymbolic/mapper/stringifier.py b/pymbolic/mapper/stringifier.py
diff --git a/pymbolic/mapper/substitutor.py b/pymbolic/mapper/substitutor.py
diff --git a/pymbolic/parser.py b/pymbolic/parser.py
diff --git a/pymbolic/primitives.py b/pymbolic/primitives.py
diff --git a/pyproject.toml b/pyproject.toml
diff --git a/test/test_pymbolic.py b/test/test_pymbolic.py
