[mypyc] Add LoadAddress op for PySet_Type & PyFrozenSet_Type

ichard26 · ichard26 · commit 7811f085d008 · 2022-07-03T15:51:14.000-04:00
This also fixes mypyc/mypyc#917 RE above, the root issue is that mypyc didn't know builtins.set was a built-in name, so it guessed it comes from the module globals. This didn't blow up anything up somehow... until the dataclasses commit[^1] which made the `__annotations__` logic for dataclasses try to better preserve the type annotations (previously they would be erased to builtins.type). This new logic would use `load_type` to load `builtins.set` (so it can be put in `__annotations__`) whick went poorly as only types registered with `load_address_op` are considered built-ins. [^1]: 1bcfc04
diff --git a/mypyc/primitives/set_ops.py b/mypyc/primitives/set_ops.py
@@ -1,13 +1,27 @@
 """Primitive set (and frozenset) ops."""
 
-from mypyc.primitives.registry import function_op, method_op, binary_op, ERR_NEG_INT
+from mypyc.primitives.registry import (
+    load_address_op, function_op, method_op, binary_op, ERR_NEG_INT
+)
 from mypyc.ir.ops import ERR_MAGIC, ERR_FALSE
 from mypyc.ir.rtypes import (
     object_rprimitive, bool_rprimitive, set_rprimitive, c_int_rprimitive, pointer_rprimitive,
     bit_rprimitive
 )
 
 
+# Get the 'builtins.set' type object.
+load_address_op(
+    name='builtins.set',
+    type=object_rprimitive,
+    src='PySet_Type')
+
+# Get the 'builtins.frozenset' tyoe object.
+load_address_op(
+    name='builtins.frozenset',
+    type=object_rprimitive,
+    src='PyFrozenSet_Type')
+
 # Construct an empty set.
 new_set_op = function_op(
     name='builtins.set',
diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
@@ -3,7 +3,7 @@
 
 from typing import (
     TypeVar, Generic, List, Iterator, Iterable, Dict, Optional, Tuple, Any, Set,
-    overload, Mapping, Union, Callable, Sequence,
+    overload, Mapping, Union, Callable, Sequence, FrozenSet
 )
 
 T = TypeVar('T')
@@ -213,6 +213,12 @@ def pop(self) -> T: pass
     def update(self, x: Iterable[S]) -> None: pass
     def __or__(self, s: Set[S]) -> Set[Union[T, S]]: ...
 
+class frozenset(Generic[T]):
+    def __init__(self, i: Optional[Iterable[T]] = None) -> None: pass
+    def __iter__(self) -> Iterator[T]: pass
+    def __len__(self) -> int: pass
+    def __or__(self, s: FrozenSet[S]) -> FrozenSet[Union[T, S]]: ...
+
 class slice: pass
 
 class range(Iterable[int]):
diff --git a/mypyc/test-data/fixtures/typing-full.pyi b/mypyc/test-data/fixtures/typing-full.pyi
@@ -19,6 +19,8 @@ TypeVar = 0
 Generic = 0
 Protocol = 0
 Tuple = 0
+Set = 0
+FrozenSet = 0
 Callable = 0
 _promote = 0
 NamedTuple = 0
diff --git a/mypyc/test-data/run-python37.test b/mypyc/test-data/run-python37.test
@@ -3,7 +3,7 @@
 [case testRunDataclass]
 import dataclasses
 from dataclasses import dataclass, field
-from typing import Set, List, Callable, Any
+from typing import Set, FrozenSet, List, Callable, Any
 
 @dataclass
 class Person1:
@@ -68,8 +68,13 @@ class Person4:
     def name(self) -> str:
         return self._name
 
+@dataclass
+class Person5:
+    friends: Set[str] = field(default_factory=set)
+    parents: FrozenSet[str] = frozenset()
+
 [file other.py]
-from native import Person1, Person1b, Person2, Person3, Person4, testBool
+from native import Person1, Person1b, Person2, Person3, Person4, Person5, testBool
 i1 = Person1(age = 5, name = 'robot')
 assert i1.age == 5
 assert i1.name == 'robot'
@@ -117,6 +122,7 @@ assert i8 > i9
 
 assert Person1.__annotations__ == {'age': int, 'name': str}
 assert Person2.__annotations__ == {'age': int, 'name': str}
+assert Person5.__annotations__ == {'friends': set, 'parents': frozenset}
 
 [file driver.py]
 import sys
