make index typing less strict (#42)

* make index typing less strict

* IList is Sequence so just use Sequence check

* make more generic

* Adding test to check error

Author: weinbe58

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

@@ -1,12 +1,14 @@
 from ._dialect import dialect as dialect
 from ._interface import (
+    col_ypos as col_ypos,
     from_positions as from_positions,
     get as get,
     get_xpos as get_xpos,
     get_ypos as get_ypos,
     new as new,
     positions as positions,
     repeat as repeat,
+    row_xpos as row_xpos,
     scale as scale,
     shape as shape,
     shift as shift,

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

@@ -22,10 +22,12 @@ def get_indices(size: int, index: Any) -> ilist.IList[int, Any]:
             raise IndexError("Index out of range")
 
         return ilist.IList([index])
-    elif isinstance(index, ilist.IList):
-        return index
-    else:
-        raise TypeError("Index must be an int, slice, or IList")
+
+    index = ilist.IList(list(index))
+    if any(not isinstance(i, int) for i in index.data):
+        raise TypeError("Index must be an int, slice, or Sequence of ints")
+
+    return index
 
 
 @dataclasses.dataclass

test/grid/test_types.py

@@ -219,6 +219,10 @@ def test_shift_subgrid_y(self, y_indices, y_shift, expected_grid):
         shifted_grid = self.grid_obj.shift_subgrid_y(y_indices, y_shift)
         assert shifted_grid.is_equal(expected_grid)
 
+    def test_invalid_slice(self):
+        with pytest.raises(TypeError):
+            self.grid_obj[[1.5], ilist.IList([0, 1])]  # type: ignore
+
     def test_scale(self):
         scaled_grid = self.grid_obj.scale(2, 3)
         expected_grid = Grid(