Ensure that slicing works on TupleOf as well as ListOf in the compiler.

Author: braxtonmckee

typed_python/compiler/tests/tuple_of_compilation_test.py

@@ -218,3 +218,17 @@ def makeT(aList: ListOf(int)):
             return TupleOf(int)(aList)
 
         assert len(makeT([])) == 0
+
+    def test_slice_list_of_compiled(self):
+        @Entrypoint
+        def sliceIt(aLst: ListOf(int), x: int):
+            return aLst[:x]
+
+        assert sliceIt((1, 2, 3), 1) == [1]
+
+    def test_slice_tuple_of_compiled(self):
+        @Entrypoint
+        def sliceIt(aTup: TupleOf(int), x: int):
+            return aTup[:x]
+
+        assert sliceIt((1, 2, 3), 1) == (1,)

typed_python/compiler/type_wrappers/list_of_wrapper.py

@@ -37,85 +37,6 @@ def list_of_extend(aList, extendWith):
             aList.append(thing)
 
 
-def min(a, b):
-    return a if a < b else b
-
-
-def max(a, b):
-    return b if a < b else a
-
-
-def list_of_slice(aList, start, stop, step):
-    if step is None:
-        return list_of_slice(aList, start, stop, 1)
-    if step == 0:
-        raise ValueError("slice step cannot be zero")
-    if step > 0:
-        if start is None:
-            return list_of_slice(aList, 0, stop, step)
-
-        if stop is None:
-            return list_of_slice(aList, start, len(aList), step)
-
-        if start < 0:
-            start += len(aList)
-        if stop < 0:
-            stop += len(aList)
-
-        start = max(0, start)
-        stop = min(len(aList), stop)
-
-        outList = type(aList)()
-
-        count = max(0, (stop - start + (step - 1)) // step)
-
-        outList.reserve(count)
-        outList.setSizeUnsafe(count)
-
-        writeP = outList.pointerUnsafe(0)
-        readP = aList.pointerUnsafe(start)
-
-        i = start
-        while i < stop:
-            writeP.initialize(readP.get())
-            writeP += 1
-            i += step
-            readP += step
-
-        return outList
-    else:
-        if start is None:
-            return list_of_slice(aList, len(aList) - 1, stop, step)
-        if stop is None:
-            return list_of_slice(aList, start, -1 - len(aList), step)
-
-        if start < 0:
-            start += len(aList)
-        if stop < 0:
-            stop += len(aList)
-
-        start = min(len(aList) - 1, start)
-        stop = max(-1, stop)
-
-        outList = type(aList)()
-        count = max(0, (start - stop + (-step - 1)) // (-step))
-
-        outList.reserve(count)
-        outList.setSizeUnsafe(count)
-
-        writeP = outList.pointerUnsafe(0)
-        readP = aList.pointerUnsafe(start)
-
-        i = start
-        while i > stop:
-            writeP.initialize(readP.get())
-            writeP += 1
-            i += step
-            readP += step
-
-        return outList
-
-
 class ListOfWrapper(TupleOrListOfWrapper):
     is_pod = False
     is_empty = False
@@ -448,18 +369,6 @@ def createEmptyList(self, context, out):
             )  # data
         )
 
-    def convert_getslice(self, context, instance, lower, upper, step):
-        return context.call_py_function(
-            list_of_slice,
-            (
-                instance,
-                lower or context.constant(None),
-                upper or context.constant(None),
-                step or context.constant(None)
-            ),
-            {}
-        )
-
     def convert_setitem(self, context, expr, index, item):
         if item is None:
             return None

typed_python/compiler/type_wrappers/tuple_of_wrapper.py

@@ -110,6 +110,14 @@ def tuple_of_hash(instance):
     return val
 
 
+def min(a, b):
+    return a if a < b else b
+
+
+def max(a, b):
+    return b if a < b else a
+
+
 class PreReservedTupleOrList(CompilableBuiltin):
     def __init__(self, tupleType):
         super().__init__()
@@ -210,6 +218,75 @@ def concatenate_tuple_or_list(l, r):
     return result
 
 
+def list_or_tupleof_of_slice(aList, start, stop, step):
+    if step is None:
+        return list_or_tupleof_of_slice(aList, start, stop, 1)
+    if step == 0:
+        raise ValueError("slice step cannot be zero")
+    if step > 0:
+        if start is None:
+            return list_or_tupleof_of_slice(aList, 0, stop, step)
+
+        if stop is None:
+            return list_or_tupleof_of_slice(aList, start, len(aList), step)
+
+        if start < 0:
+            start += len(aList)
+        if stop < 0:
+            stop += len(aList)
+
+        start = max(0, start)
+        stop = min(len(aList), stop)
+
+        count = max(0, (stop - start + (step - 1)) // step)
+
+        outLstOrTup = PreReservedTupleOrList(type(aList))(count)
+
+        outLstOrTup.setSizeUnsafe(count)
+
+        writeP = outLstOrTup.pointerUnsafe(0)
+        readP = aList.pointerUnsafe(start)
+
+        i = start
+        while i < stop:
+            writeP.initialize(readP.get())
+            writeP += 1
+            i += step
+            readP += step
+
+        return outLstOrTup
+    else:
+        if start is None:
+            return list_or_tupleof_of_slice(aList, len(aList) - 1, stop, step)
+        if stop is None:
+            return list_or_tupleof_of_slice(aList, start, -1 - len(aList), step)
+
+        if start < 0:
+            start += len(aList)
+        if stop < 0:
+            stop += len(aList)
+
+        start = min(len(aList) - 1, start)
+        stop = max(-1, stop)
+
+        count = max(0, (start - stop + (-step - 1)) // (-step))
+
+        outLstOrTup = PreReservedTupleOrList(type(aList))(count)
+        outLstOrTup.setSizeUnsafe(count)
+
+        writeP = outLstOrTup.pointerUnsafe(0)
+        readP = aList.pointerUnsafe(start)
+
+        i = start
+        while i > stop:
+            writeP.initialize(readP.get())
+            writeP += 1
+            i += step
+            readP += step
+
+        return outLstOrTup
+
+
 class TupleOrListOfWrapper(RefcountedWrapper):
     is_pod = False
     is_empty = False
@@ -364,6 +441,18 @@ def convert_set_attribute(self, context, expr, attr, val):
     def convert_hash(self, context, expr):
         return context.call_py_function(tuple_of_hash, (expr,), {})
 
+    def convert_getslice(self, context, instance, lower, upper, step):
+        return context.call_py_function(
+            list_or_tupleof_of_slice,
+            (
+                instance,
+                lower or context.constant(None),
+                upper or context.constant(None),
+                step or context.constant(None)
+            ),
+            {}
+        )
+
     def convert_getitem(self, context, expr, item):
         if item is None or expr is None:
             return None