Skip to content

Commit 82de0d8

Browse files
authored
[mypyc] Fixing iteration over NamedTuple objects. (#18254)
Fixes mypyc/mypyc#1063. Adding elif-blocks for getting RType of items in classes derived from `NamedTuple`.
1 parent 802266b commit 82de0d8

File tree

2 files changed

+43
-5
lines changed

2 files changed

+43
-5
lines changed

mypyc/irbuild/builder.py

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
Type,
5353
TypedDictType,
5454
TypeOfAny,
55+
TypeVarLikeType,
5556
UninhabitedType,
5657
UnionType,
5758
get_proper_type,
@@ -926,11 +927,22 @@ def get_sequence_type_from_type(self, target_type: Type) -> RType:
926927
return RUnion.make_simplified_union(
927928
[self.get_sequence_type_from_type(item) for item in target_type.items]
928929
)
929-
assert isinstance(target_type, Instance), target_type
930-
if target_type.type.fullname == "builtins.str":
931-
return str_rprimitive
932-
else:
933-
return self.type_to_rtype(target_type.args[0])
930+
elif isinstance(target_type, Instance):
931+
if target_type.type.fullname == "builtins.str":
932+
return str_rprimitive
933+
else:
934+
return self.type_to_rtype(target_type.args[0])
935+
# This elif-blocks are needed for iterating over classes derived from NamedTuple.
936+
elif isinstance(target_type, TypeVarLikeType):
937+
return self.get_sequence_type_from_type(target_type.upper_bound)
938+
elif isinstance(target_type, TupleType):
939+
# Tuple might have elements of different types.
940+
rtypes = {self.mapper.type_to_rtype(item) for item in target_type.items}
941+
if len(rtypes) == 1:
942+
return rtypes.pop()
943+
else:
944+
return RUnion.make_simplified_union(list(rtypes))
945+
assert False, target_type
934946

935947
def get_dict_base_type(self, expr: Expression) -> list[Instance]:
936948
"""Find dict type of a dict-like expression.

mypyc/test-data/run-loops.test

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -545,3 +545,29 @@ def test_range_object() -> None:
545545
r4 = range(4, 12, 0)
546546
except ValueError as e:
547547
assert "range() arg 3 must not be zero" in str(e)
548+
549+
[case testNamedTupleLoop]
550+
from collections.abc import Iterable
551+
from typing import NamedTuple, Any
552+
from typing_extensions import Self
553+
554+
555+
class Vector2(NamedTuple):
556+
x: int
557+
y: float
558+
559+
@classmethod
560+
def from_iter(cls, iterable: Iterable[Any]) -> Self:
561+
return cls(*iter(iterable))
562+
563+
def __neg__(self) -> Self:
564+
return self.from_iter(-c for c in self)
565+
566+
[file driver.py]
567+
import native
568+
print(-native.Vector2(2, -3.1))
569+
print([x for x in native.Vector2(4, -5.2)])
570+
571+
[out]
572+
Vector2(x=-2, y=3.1)
573+
\[4, -5.2]

0 commit comments

Comments
 (0)