Make sure we can compile expressions like 'x, y = OneOf(None, Tuple(int, int))((1, 2))'.

Author: braxtonmckee

typed_python/compiler/tests/one_of_compilation_test.py

@@ -724,3 +724,11 @@ def toString(x: OneOf("hi", "bye")):  # noqa
             return f"its: {x}"
 
         toString('hi')
+
+    def test_unpack_oneof_none_or_tuple(self):
+        @Entrypoint
+        def unpackIt(x: OneOf(None, Tuple(int, int))):
+            a, b = x
+            return a + b
+
+        unpackIt((1, 2))

typed_python/compiler/tests/tuple_compilation_test.py

@@ -398,3 +398,18 @@ def checkIn():
             return None
 
         assert checkIn.resultTypeFor().typeRepresentation == int
+
+    def test_iterate_tuple(self):
+        T = Tuple(int, float, str)
+
+        @Entrypoint
+        def iterateIt(x: T):
+            res = ListOf(OneOf(int, float, str))()
+
+            for e in x:
+                res.append(e)
+
+            return res
+
+        tup = T((1, 2, '3'))
+        assert iterateIt(tup) == ListOf(OneOf(int, float, str))(tup)

typed_python/compiler/type_wrappers/none_wrapper.py

@@ -29,6 +29,9 @@ def __init__(self):
     def getCompileTimeConstant(self):
         return None
 
+    def isIterable(self):
+        return False
+
     def convert_default_initialize(self, context, target):
         pass
 

typed_python/compiler/type_wrappers/one_of_wrapper.py

@@ -57,6 +57,34 @@ def getNativeLayoutType(self):
     def convert_which_native(self, expr):
         return expr.ElementPtrIntegers(0, 0).load()
 
+    def get_iteration_expressions(self, context, expr):
+        itExprs = None
+
+        for ix in range(len(self.typeRepresentation.Types)):
+            T = self.typeRepresentation.Types[ix]
+
+            if typeWrapper(T).isIterable == "Maybe":
+                return
+            elif typeWrapper(T).isIterable is False:
+                pass
+            else:
+                # we don't know how to union two sets of iteration expressions
+                if itExprs is not None:
+                    return None
+
+                itExprs = typeWrapper(T).get_iteration_expressions(context, expr.refAs(ix))
+                itExprIx = ix
+
+        if itExprs is not None:
+            with context.ifelse(self.convert_which_native(expr.expr).eq(itExprIx)) as (ifTrue, ifFalse):
+                with ifFalse:
+                    context.pushException(
+                        TypeError,
+                        "Instance of type {self.typeRepresentation.__name__} is not iterable"
+                    )
+
+            return itExprs
+
     def unwrap(self, context, expr, generator):
         """Call 'generator' on 'expr' cast down to each subtype and combine the results.
         """

typed_python/compiler/type_wrappers/tuple_of_wrapper.py

@@ -318,6 +318,9 @@ def __init__(self, t):
             ('data', native_ast.UInt8Ptr)
         ), name='TupleOfLayout' if self.is_tuple else 'ListOfLayout').pointer()
 
+    def isIterable(self):
+        return True
+
     def has_fastnext_iter(self):
         """If we call '__iter__' on instances of this type, will they support __fastnext__?"""
         return True

typed_python/compiler/type_wrappers/tuple_wrapper.py

@@ -15,7 +15,8 @@
 from typed_python.compiler.type_wrappers.wrapper import Wrapper
 from typed_python.compiler.merge_type_wrappers import mergeTypes
 from typed_python import (
-    _types, Int32, Tuple, NamedTuple, Function, Dict, Set, ConstDict, ListOf, TupleOf
+    _types, Int32, Tuple, NamedTuple, Function, Dict, Set, ConstDict, ListOf,
+    TupleOf, PointerTo, pointerTo, TypeFunction, OneOf, Class, Member, Final
 )
 import typed_python._types
 from typed_python.compiler.conversion_level import ConversionLevel
@@ -195,6 +196,9 @@ def __init__(self, t):
         self._is_pod = all(typeWrapper(possibility).is_pod for possibility in self.subTypeWrappers)
         self.is_default_constructible = _types.is_default_constructible(t)
 
+    def isIterable(self):
+        return True
+
     @property
     def unionType(self):
         if self._unionType is None and self.typeRepresentation.ElementTypes:
@@ -220,6 +224,23 @@ def is_pod(self):
     def getNativeLayoutType(self):
         return self.layoutType
 
+    def convert_attribute(self, context, instance, attribute):
+        if attribute in ["__iter__"]:
+            return instance.changeType(BoundMethodWrapper.Make(self, attribute))
+
+        return super().convert_attribute(context, instance, attribute)
+
+    def convert_method_call(self, context, instance, methodname, args, kwargs):
+        if methodname == '__iter__' and not args:
+            return typeWrapper(TupleIterator(self.typeRepresentation)).convert_type_call(
+                context,
+                None,
+                [],
+                dict(pos=context.constant(-1), tup=instance)
+            )
+
+        return super().convert_method_call(context, instance, methodname, args, kwargs)
+
     def convert_initialize_from_args(self, context, target, *args):
         assert len(args) <= len(self.byteOffsets)
 
@@ -608,6 +629,9 @@ def __init__(self, t):
         self.namesToIndices = {n: i for i, n in enumerate(t.ElementNames)}
         self.namesToTypes = {n: t.ElementTypes[i] for i, n in enumerate(t.ElementNames)}
 
+    def isIterable(self):
+        return True
+
     def has_fastnext_iter(self):
         if self.isSubclassOfNamedTuple:
             return "__iter__" in self.typeRepresentation.__dict__
@@ -869,3 +893,24 @@ def convert_bin_op_reverse(self, context, r, op, l, inplace):
             return self.convert_method_call(context, r, magic, (l,), {})
 
         return super().convert_bin_op_reverse(context, r, op, l, inplace)
+
+
+@TypeFunction
+def TupleIterator(T):
+    EltT = OneOf(*T.ElementTypes)
+
+    class TupleIterator(Class, Final, __name__=f"TupleIterator({T.__name__})"):
+        pos = Member(int, nonempty=True)
+        elt = Member(EltT)
+        tup = Member(T, nonempty=True)
+
+        def __fastnext__(self):
+            self.pos += 1
+
+            if self.pos < len(self.tup):
+                self.elt = self.tup[self.pos]
+                return pointerTo(self).elt
+            else:
+                return PointerTo(EltT)()
+
+    return TupleIterator