mypy fixes

inducer · Nov 11, 2024 · 5e5be1d · 5e5be1d
1 parent 0ef534e
commit 5e5be1d
Showing 1 changed file with 4 additions and 20 deletions.
diff --git a/loopy/statistics.py b/loopy/statistics.py
@@ -196,7 +196,7 @@ def __repr__(self):
 # {{{ ToCountMap
 
 Countable = Union["Op", "MemAccess", "Sync"]
-CountT = TypeVar("CountT")
+CountT = TypeVar("CountT", int, GuardedPwQPolynomial)
 
 
 class ToCountMap(Generic[CountT]):
@@ -966,7 +966,7 @@ def new_poly_map(self, count_map) -> ToCountPolynomialMap:
     def _new_zero_map(self) -> ToCountPolynomialMap:
         return self.new_poly_map({})
 
-    def combine(self, values: Iterable[ToCountMap]) -> ToCountPolynomialMap:
+    def combine(self, values: Iterable[ToCountPolynomialMap]) -> ToCountPolynomialMap:
         return sum(values, self._new_zero_map())
 
     def map_tagged_expression(
@@ -975,7 +975,7 @@ def map_tagged_expression(
         return self.rec(expr.expr, expr.tags)
 
     def map_constant(
-            self, expr: Expression, tags: frozenset[Tag]
+            self, expr: object, tags: frozenset[Tag]
             ) -> ToCountPolynomialMap:
         return self._new_zero_map()
 
@@ -1027,14 +1027,6 @@ def map_if(
         return self.rec(expr.condition, tags) + self.rec(expr.then, tags) \
                + self.rec(expr.else_, tags)
 
-    def map_if_positive(
-            self, expr: p.IfPositive, tags: frozenset[Tag]) -> ToCountMap:
-        warn_with_kernel(self.knl, "summing_if_branches",
-                         "%s counting sum of if-expression branches."
-                         % type(self).__name__)
-        return self.rec(expr.criterion, tags) + self.rec(expr.then, tags) \
-               + self.rec(expr.else_, tags)
-
     def map_common_subexpression(
             self, expr: p.CommonSubexpression, tags: frozenset[Tag]
             ) -> ToCountPolynomialMap:
@@ -1204,14 +1196,6 @@ def map_if(self, expr: p.If, tags: frozenset[Tag]) -> ToCountPolynomialMap:
         return self.rec(expr.condition, tags) + self.rec(expr.then, tags) \
                + self.rec(expr.else_, tags)
 
-    def map_if_positive(
-            self, expr: p.IfPositive, tags: frozenset[Tag]) -> ToCountPolynomialMap:
-        warn_with_kernel(self.knl, "summing_ifpos_branches_ops",
-                         "ExpressionOpCounter counting ops as sum of "
-                         "if_pos-statement branches.")
-        return self.rec(expr.criterion, tags) + self.rec(expr.then, tags) \
-               + self.rec(expr.else_, tags)
-
     def map_min(
             self, expr: Union[p. Min, p.Max], tags: frozenset[Tag]
             ) -> ToCountPolynomialMap:
@@ -2418,7 +2402,7 @@ def gather_access_footprint_bytes(
     """
 
     from loopy.preprocess import infer_unknown_types
-    kernel = infer_unknown_types(t_unit, expect_completion=True)
+    t_unit = infer_unknown_types(t_unit, expect_completion=True)
 
     fp = gather_access_footprints(t_unit, ignore_uncountable=ignore_uncountable)