[IR] introduce slice support (#2302)

Author: Johansmm

onnxscript/ir/_core.py

@@ -2282,7 +2282,12 @@ def doc_string(self, value: str | None) -> None:
     def opset_imports(self) -> dict[str, int]:
         return self._opset_imports
 
-    def __getitem__(self, index: int) -> Node:
+    @typing.overload
+    def __getitem__(self, index: int) -> Node: ...
+    @typing.overload
+    def __getitem__(self, index: slice) -> tuple[Node, ...]: ...
+
+    def __getitem__(self, index):
         return self._nodes[index]
 
     def __len__(self) -> int:
@@ -2712,7 +2717,12 @@ def __init__(
         self._metadata_props: dict[str, str] | None = metadata_props
         self._nodes: tuple[Node, ...] = tuple(nodes)
 
-    def __getitem__(self, index: int) -> Node:
+    @typing.overload
+    def __getitem__(self, index: int) -> Node: ...
+    @typing.overload
+    def __getitem__(self, index: slice) -> tuple[Node, ...]: ...
+
+    def __getitem__(self, index):
         return self._nodes[index]
 
     def __len__(self) -> int:
@@ -2961,7 +2971,12 @@ def outputs(self) -> MutableSequence[Value]:
     def attributes(self) -> OrderedDict[str, Attr]:
         return self._attributes
 
-    def __getitem__(self, index: int) -> Node:
+    @typing.overload
+    def __getitem__(self, index: int) -> Node: ...
+    @typing.overload
+    def __getitem__(self, index: slice) -> tuple[Node, ...]: ...
+
+    def __getitem__(self, index):
         return self._graph.__getitem__(index)
 
     def __len__(self) -> int:

onnxscript/ir/_linked_list.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from typing import Generic, Iterable, Iterator, Sequence, TypeVar
+from typing import Generic, Iterable, Iterator, Sequence, TypeVar, overload
 
 T = TypeVar("T")
 
@@ -136,27 +136,54 @@ def __len__(self) -> int:
         )
         return self._length
 
-    def __getitem__(self, index: int) -> T:
+    @overload
+    def __getitem__(self, index: int) -> T: ...
+    @overload
+    def __getitem__(self, index: slice) -> list[T]: ...
+
+    def __getitem__(self, index):
         """Get the node at the given index.
 
         Complexity is O(n).
         """
-        if index >= self._length or index < -self._length:
+        if isinstance(index, slice):
+            start, stop, step = index.indices(self._length)
+        elif index < 0:
+            start, stop, step = index, -self._length - 1, -1
+        else:
+            start, stop, step = index, self._length, 1
+
+        if start >= self._length or start < -self._length:
             raise IndexError(
                 f"Index out of range: {index} not in range [-{self._length}, {self._length})"
             )
-        if index < 0:
+        if (step < 0 and stop >= start) or (step > 0 and stop <= start):
+            return []
+        if step < 0:
             # Look up from the end of the list
             iterator = reversed(self)
             item = next(iterator)
-            for _ in range(-index - 1):
+            # Skip index to match with start
+            for _ in range(-start - 1 if start < 0 else self._length - start - 1):
                 item = next(iterator)
         else:
             iterator = iter(self)  # type: ignore[assignment]
             item = next(iterator)
-            for _ in range(index):
+            # Skip index to match with start
+            for _ in range(start):
                 item = next(iterator)
-        return item
+
+        # Return a single element if index is an integer
+        if isinstance(index, int):
+            return item
+
+        # Return a list if index is a slice
+        items = [item]
+        for i in range(1, abs(stop - start)):
+            item = next(iterator)
+            if i % step == 0:
+                items.append(item)
+        return items
 
     def _insert_one_after(
         self,

onnxscript/ir/_linked_list_test.py

@@ -373,6 +373,15 @@ def test_insert_after_supports_taking_elements_from_another_doubly_linked_list(
         self.assertEqual(len(other_linked_list), 1)
         self.assertEqual([elem.value for elem in other_linked_list], [42])
 
+    @parameterized.parameterized.expand(
+        [(s, t, p) for s in [-2, 0, 2, 3] for t in [2, -1, -2] for p in [-3, -2, -1, 1, 2]]
+    )
+    def test_get_item_slice(self, start, stop, step):
+        elems = [_TestElement(i) for i in range(5)]
+        linked_list = _linked_list.DoublyLinkedSet(elems)
+        self.assertEqual(len(linked_list), 5)
+        self.assertEqual(linked_list[start:stop:step], elems[start:stop:step])
+
 
 if __name__ == "__main__":
     unittest.main()