update type stub to include indexing with np.integer (#36)

quantumlib · Feb 12, 2025 · 05d5eb4 · 05d5eb4
2 parents 27f1bed + 74b45c5
commit 05d5eb4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/test/test_value_array.py b/test/test_value_array.py
@@ -153,6 +153,9 @@ def test_multi_index() -> None:
     with pytest.raises(IndexError):
         _ = (m * [[2, 3, 4], [5, 6, 7], [8, 9, 10]])[1:3, 25483]
 
+    idx = np.argmax([1, 4, -1])
+    assert (m * [1, 2, 4])[idx] == m * 2
+
 
 def test_predicate_index() -> None:
     from tunits.units import m

diff --git a/tunits/core/__init__.pyi b/tunits/core/__init__.pyi
@@ -366,7 +366,7 @@ class ValueArray(Generic[ValueType2], WithUnit):
     @overload
     def __getitem__(self, key: str | Value) -> NDArray[Any]: ...
     @overload
-    def __getitem__(self, key: int) -> ValueType2: ...
+    def __getitem__(self, key: int | np.integer) -> ValueType2: ...
     @overload
     def __getitem__(self: ArrayType, key: slice | tuple[Any, ...] | NDArray[Any]) -> ArrayType: ...
     def value_in_base_units(self) -> NDArray[Any]: ...