From e86a15e39c65932773b2fee37d6de0aebcaacf52 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Wed, 28 Oct 2026 14:45:20 -0500 Subject: [PATCH 1/3] Typing: ExpansionState.instruction can be None --- loopy/symbolic.py | 6 ++++-- loopy/transform/iname.py | 17 +++++++++++------ loopy/transform/loop_fusion.py | 1 + loopy/transform/precompute.py | 1 + 4 files changed, 17 insertions(+), 8 deletions(-) diff --git a/loopy/symbolic.py b/loopy/symbolic.py index 01d09fe7b..a9f074be8 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -1214,7 +1214,7 @@ class ExpansionState: """ kernel: LoopKernel - instruction: InstructionBase + instruction: InstructionBase | None stack: tuple[ConcreteMatchable, ...] """A tuple representing the current expansion stack, as a tuple of ``(name, tag)`` pairs. @@ -1238,6 +1238,7 @@ def __hash__(self) -> int: @property def insn_id(self) -> str: + assert self.instruction is not None return self.instruction.id def apply_arg_context(self, expr: Expression) -> Expression: @@ -1530,7 +1531,7 @@ def map_subst_rule( def __call__(self, # pyright: ignore[reportIncompatibleMethodOverride] expr: Expression, kernel: LoopKernel, - insn: InstructionBase, + insn: InstructionBase | None, *args: P.args, **kwargs: P.kwargs) -> Expression: """ :arg insn: A :class:`~loopy.kernel.InstructionBase` of which *expr* is @@ -1691,6 +1692,7 @@ def map_subst_rule( from loopy.match import ConcreteMatchable new_stack = (*expn_state.stack, ConcreteMatchable(name, tags)) + assert expn_state.instruction is not None if self.within(expn_state.kernel, expn_state.instruction, new_stack): # expand rule = self.rules[name] diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index d4735cf71..0fb9bd7a0 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -199,6 +199,7 @@ def map_reduction(self, expr: Reduction, /, expn_state: ExpansionState) -> Expression: if (self.iname_to_split in expr.inames and self.iname_to_split not in expn_state.arg_context + and expn_state.instruction is not None and self.within( expn_state.kernel, expn_state.instruction)): @@ -217,6 +218,7 @@ def map_variable(self, expr: prim.Variable, /, expn_state: ExpansionState) -> Expression: if (expr.name == self.iname_to_split and self.iname_to_split not in expn_state.arg_context + and expn_state.instruction is not None and self.within( expn_state.kernel, expn_state.instruction)): @@ -875,6 +877,7 @@ def __init__(self, def map_reduction(self, expr: Reduction, /, expn_state: ExpansionState) -> Expression: if (set(expr.inames) & self.old_inames_set + and expn_state.instruction is not None and self.within( expn_state.kernel, expn_state.instruction, @@ -896,15 +899,16 @@ def map_variable(self, expr: prim.Variable, /, expn_state: ExpansionState) -> Expression: new_name = self.old_to_new.get(expr.name) - if (new_name is None - or expr.name in expn_state.arg_context - or not self.within( + if (new_name is not None + and expr.name not in expn_state.arg_context + and expn_state.instruction is not None + and self.within( expn_state.kernel, expn_state.instruction, expn_state.stack)): - return super().map_variable(expr, expn_state) - else: return prim.Variable(new_name) + else: + return super().map_variable(expr, expn_state) @override def map_instruction(self, @@ -1346,6 +1350,7 @@ def map_reduction(self, expr: Reduction, /, raise NotImplementedError() if (self.inames <= set(expr.inames) + and expn_state.instruction is not None and self.within( expn_state.kernel, expn_state.instruction, @@ -1544,7 +1549,7 @@ def get_cache_key(self, expr: Expression, expn_state: ExpansionState) -> Hashabl @override def map_reduction(self, expr: Reduction, /, expn_state: ExpansionState) -> Expression: - within = self.within( + within = expn_state.instruction is not None and self.within( expn_state.kernel, expn_state.instruction, expn_state.stack) diff --git a/loopy/transform/loop_fusion.py b/loopy/transform/loop_fusion.py index 99dbf033f..6ec2cbe4a 100644 --- a/loopy/transform/loop_fusion.py +++ b/loopy/transform/loop_fusion.py @@ -677,6 +677,7 @@ def map_reduction(self, outer_redn_inames: frozenset[InameStr] = _EMPTY_INAME_SET ): redn_inames = frozenset(expr.inames) + assert expn_state.instruction is not None iname_chain = ( expn_state.instruction.within_inames | outer_redn_inames | redn_inames ) diff --git a/loopy/transform/precompute.py b/loopy/transform/precompute.py index 3a638e4c7..db89f469d 100644 --- a/loopy/transform/precompute.py +++ b/loopy/transform/precompute.py @@ -280,6 +280,7 @@ def map_subst_rule(self, expn_state: ExpansionState): if not ( name == self.subst_name + and expn_state.instruction is not None and self.within( expn_state.kernel, expn_state.instruction, From 50cf18654b19a178f04e537bbaf7e8971fea812b Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Wed, 28 Oct 2026 14:45:25 -0500 Subject: [PATCH 2/3] Update baseline --- .basedpyright/baseline.json | 152 ------------------------------------ 1 file changed, 152 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 4dacf996f..309719deb 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -21195,14 +21195,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 16, - "endColumn": 31, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -21243,14 +21235,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 16, - "endColumn": 31, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -21491,14 +21475,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 57, - "endColumn": 72, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -21555,14 +21531,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 16, - "endColumn": 31, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -38643,94 +38611,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 8, - "endColumn": 23, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 30, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 30, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 47, - "endColumn": 62, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 47, - "endColumn": 62, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 29, - "endColumn": 50, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 41, - "endColumn": 51, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 36, - "endColumn": 31, - "lineCount": 3 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 20, - "endColumn": 37, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 55, - "endColumn": 76, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 49, - "endColumn": 67, - "lineCount": 1 - } - }, { "code": "reportUnannotatedClassAttribute", "range": { @@ -57867,14 +57747,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 66, - "endColumn": 70, - "lineCount": 1 - } - }, { "code": "reportPrivateLocalImportUsage", "range": { @@ -59541,14 +59413,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 32, - "endColumn": 36, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -72725,22 +72589,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 16, - "endColumn": 46, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 11, - "endColumn": 24, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { From 59a518efc20e4a47adf48831c5c3d76756f3b04a Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Wed, 28 Oct 2026 14:45:40 -0500 Subject: [PATCH 3/3] Add preliminary ty configuration --- pyproject.toml | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 9ac9257ed..5bcb6512d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -183,6 +183,19 @@ filterwarnings = [ 'ignore:.*:loopy.diagnostic.DirectCallUncachedWarning', ] +[tool.ty] +baseline = ".ty-baseline.json" + +[tool.ty.analysis] +return-type-inference = "always" + +# https://docs.astral.sh/ty/coming-from-mypy-or-pyright/#recommended-configuration +[tool.ty.rules] +dynamic-function-decorator-return = "error" +missing-type-argument = "error" +possibly-unresolved-reference = "warn" +unsound-return-statement = "error" + [tool.basedpyright] reportImplicitStringConcatenation = "none" reportUnnecessaryIsInstance = "none"