updating api for grid.get_view to be more flexible (#29)

Author: weinbe58

src/bloqade/geometry/dialects/grid/types.py

@@ -205,19 +205,47 @@ def get(self, idx: tuple[int, int]) -> tuple[float, float]:
     Nx = TypeVar("Nx")
     Ny = TypeVar("Ny")
 
+    @overload
     def get_view(
         self, x_indices: ilist.IList[int, Nx], y_indices: ilist.IList[int, Ny]
-    ) -> "Grid[Nx, Ny]":
+    ) -> "Grid[Nx, Ny]": ...
+
+    @overload
+    def get_view(
+        self, x_indices: Sequence[int], y_indices: ilist.IList[int, Ny]
+    ) -> "Grid[Any, Ny]": ...
+
+    @overload
+    def get_view(
+        self, x_indices: ilist.IList[int, Nx], y_indices: Sequence[int]
+    ) -> "Grid[Nx, Any]": ...
+
+    @overload
+    def get_view(
+        self, x_indices: Sequence[int], y_indices: Sequence[int]
+    ) -> "Grid[Any, Any]": ...
+
+    def get_view(self, x_indices, y_indices) -> "Grid":
         """Get a sub-grid view based on the specified x and y indices.
 
         Args:
-            x_indices (ilist.IList[int, Nx]): The x indices to include in the sub-grid.
-            y_indices (ilist.IList[int, Ny]): The y indices to include in the sub-grid.
+            x_indices (Sequence[int]): The x indices to include in the sub-grid.
+            y_indices (Sequence[int]): The y indices to include in the sub-grid.
 
         Returns:
-            Grid[Nx, Ny]: The sub-grid view.
+            Grid: The sub-grid view.
         """
-        return SubGrid(parent=self, x_indices=x_indices, y_indices=y_indices)
+        if isinstance(x_indices, ilist.IList):
+            x_indices = x_indices.data
+
+        if isinstance(y_indices, ilist.IList):
+            y_indices = y_indices.data
+
+        return SubGrid(
+            parent=self,
+            x_indices=ilist.IList(x_indices),
+            y_indices=ilist.IList(y_indices),
+        )
 
     @overload
     def __getitem__(
@@ -428,9 +456,7 @@ def __post_init__(self):
             types.Literal(len(self.y_indices)),
         )
 
-    def get_view(
-        self, x_indices: ilist.IList[int, Any], y_indices: ilist.IList[int, Any]
-    ):
+    def get_view(self, x_indices, y_indices):
         return self.parent.get_view(
             x_indices=ilist.IList([self.x_indices[x_index] for x_index in x_indices]),
             y_indices=ilist.IList([self.y_indices[y_index] for y_index in y_indices]),

test/grid/test_typeinfer.py

@@ -13,7 +13,7 @@ def test_typeinfer_1():
     def test_1(spacing: ilist.IList[float, Literal[2]]):
         return grid.new(spacing, [1.0, 2.0], 0.0, 0.0)
 
-    assert test_1.return_type.is_equal(
+    assert test_1.return_type.is_subseteq(
         grid.GridType[types.Literal(3), types.Literal(3)]
     )
 
@@ -23,7 +23,7 @@ def test_typeinfer_2():
     def test_2(spacing: ilist.IList[float, Any]):
         return grid.new(spacing, [1.0, 2.0], 0.0, 0.0)
 
-    assert test_2.return_type.is_equal(grid.GridType[types.Any, types.Literal(3)])
+    assert test_2.return_type.is_subseteq(grid.GridType[types.Any, types.Literal(3)])
 
 
 def test_typeinfer_get_index_1():