Merge pull request #5517 from esc/typed-list/extend_refine_tuple

Fix numba.typed.List extend for singleton and empty iterable

numba/tests/test_typedlist.py

@@ -738,6 +738,33 @@ def impl():
         got = impl()
         self.assertEqual(expected, got)
 
+    def test_extend_single_value_container(self):
+        @njit
+        def impl():
+            l = List()
+            l.extend((100,))
+            return l
+
+        expected = impl.py_func()
+        got = impl()
+        self.assertEqual(expected, got)
+
+    def test_extend_empty_unrefined(self):
+        # Extending an unrefined list with an empty iterable doesn't work in a
+        # jit compiled function as the list remains untyped.
+        l = List()
+        l.extend(tuple())
+        self.assertEqual(len(l), 0)
+        self.assertFalse(l._typed)
+
+    def test_extend_empty_refiend(self):
+        # Extending a refined list with an empty iterable doesn't work in a
+        # jit compiled function as the (empty) argument can't be typed
+        l = List((1,))
+        l.extend(tuple())
+        self.assertEqual(len(l), 1)
+        self.assertTrue(l._typed)
+
 
 @njit
 def cmp(a, b):

numba/typed/typedlist.py

@@ -355,13 +355,14 @@ def pop(self, i=-1):
         return _pop(self, i)
 
     def extend(self, iterable):
+        # Empty iterable, do nothing
+        if len(iterable) == 0:
+            return self
         if not self._typed:
             # Need to get the first element of the iterable to initialise the
             # type of the list. FIXME: this may be a problem if the iterable
             # can not be sliced.
             self._initialise_list(iterable[0])
-            self.append(iterable[0])
-            return _extend(self, iterable[1:])
         return _extend(self, iterable)
 
     def remove(self, item):