Making typed.List a typing Generic

numba/typed/typedlist.py

@@ -26,6 +26,7 @@
 from numba.typed import listobject
 from numba.core.errors import TypingError, LoweringError
 from numba.core.typing.templates import Signature
+import typing as pt
 
 DEFAULT_ALLOCATED = listobject.DEFAULT_ALLOCATED
 
@@ -170,7 +171,10 @@ def _from_meminfo_ptr(ptr, listtype):
     return List(meminfo=ptr, lsttype=listtype)
 
 
-class List(MutableSequence):
+T = pt.TypeVar('T')
+
+
+class List(MutableSequence, pt.Generic[T]):
     """A typed-list usable in Numba compiled functions.
 
     Implements the MutableSequence interface.
@@ -278,7 +282,7 @@ def _initialise_list(self, item):
         lsttype = types.ListType(typeof(item))
         self._list_type, self._opaque = self._parse_arg(lsttype)
 
-    def __len__(self):
+    def __len__(self) -> int:
         if not self._typed:
             return 0
         else:
@@ -317,44 +321,52 @@ def __gt__(self, other):
     def __ge__(self, other):
         return _ge(self, other)
 
-    def append(self, item):
+    def append(self, item: T) -> None:
         if not self._typed:
             self._initialise_list(item)
         _append(self, item)
 
-    def __setitem__(self, i, item):
+    def __setitem__(self, i, item: T) -> None:
         if not self._typed:
             self._initialise_list(item)
         _setitem(self, i, item)
 
-    def __getitem__(self, i):
+    # noqa F811 comments required due to github.com/PyCQA/pyflakes/issues/592
+    # noqa E704 required to follow overload style of using ... in the same line
+    @pt.overload
+    def __getitem__(self, i: int) -> T: ... # noqa: F811, E704
+    @pt.overload
+    def __getitem__(self, i: slice) -> 'List': ... # noqa: F811, E704
+
+    def __getitem__(self, i: pt.Union[int, slice] # noqa: F811
+                    ) -> pt.Union[T, 'List']:
         if not self._typed:
             raise IndexError
         else:
             return _getitem(self, i)
 
-    def __iter__(self):
+    def __iter__(self) -> pt.Iterator[T]:
         for i in range(len(self)):
             yield self[i]
 
-    def __contains__(self, item):
+    def __contains__(self, item: T) -> bool:
         return _contains(self, item)
 
-    def __delitem__(self, i):
+    def __delitem__(self, i: pt.Union[int, slice]) -> None:
         _delitem(self, i)
 
-    def insert(self, i, item):
+    def insert(self, i: int, item: T) -> None:
         if not self._typed:
             self._initialise_list(item)
         _insert(self, i, item)
 
-    def count(self, item):
+    def count(self, item: T) -> int:
         return _count(self, item)
 
-    def pop(self, i=-1):
+    def pop(self, i: int = -1) -> T:
         return _pop(self, i)
 
-    def extend(self, iterable):
+    def extend(self, iterable: pt.Iterable[T]) -> None:
         # Empty iterable, do nothing
         if len(iterable) == 0:
             return self
@@ -365,7 +377,7 @@ def extend(self, iterable):
             self._initialise_list(iterable[0])
         return _extend(self, iterable)
 
-    def remove(self, item):
+    def remove(self, item: T) -> None:
         return _remove(self, item)
 
     def clear(self):
@@ -377,7 +389,7 @@ def reverse(self):
     def copy(self):
         return _copy(self)
 
-    def index(self, item, start=None, stop=None):
+    def index(self, item: T, start: int = None, stop: int = None) -> int:
         return _index(self, item, start, stop)
 
     def sort(self, key=None, reverse=False):