Apply suggestions from code review

Co-authored-by: Joren Hammudoglu <[email protected]>

Author: crusaderky

array_api_strict/_array_object.py

@@ -1023,7 +1023,7 @@ def __rand__(self, other: Array | int, /) -> Array:
         res = self._array.__rand__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ifloordiv__(self, other: Array | complex, /) -> Array:
+    def __ifloordiv__(self, other: Array | float, /) -> Array:
         """
         Performs the operation __ifloordiv__.
         """
@@ -1034,7 +1034,7 @@ def __ifloordiv__(self, other: Array | complex, /) -> Array:
         self._array.__ifloordiv__(other._array)
         return self
 
-    def __rfloordiv__(self, other: Array | complex, /) -> Array:
+    def __rfloordiv__(self, other: Array | float, /) -> Array:
         """
         Performs the operation __rfloordiv__.
         """
@@ -1105,7 +1105,7 @@ def __imod__(self, other: Array | complex, /) -> Array:
         self._array.__imod__(other._array)
         return self
 
-    def __rmod__(self, other: Array | complex, /) -> Array:
+    def __rmod__(self, other: Array | float, /) -> Array:
         """
         Performs the operation __rmod__.
         """

array_api_strict/_creation_functions.py

@@ -24,7 +24,7 @@ class Undef(Enum):
 
 
 @contextmanager
-def allow_array() -> Generator[None, None, None]:
+def allow_array() -> Generator[None]:
     """
     Temporarily enable Array.__array__. This is needed for np.array to parse
     list of lists of Array objects.

array_api_strict/_flags.py

@@ -348,6 +348,7 @@ def __exit__(
         exc_type: type[BaseException] | None,
         exc_value: BaseException | None,
         traceback: TracebackType | None,
+        /,
     ) -> None:
         set_array_api_strict_flags(**self.old_flags)
 

array_api_strict/_utility_functions.py

@@ -56,7 +56,7 @@ def diff(
     # currently specified.
 
     # NumPy does not support prepend=None or append=None
-    kwargs: dict[str, Any] = dict(axis=axis, n=n)
+    kwargs: dict[str, int | npt.NDArray[Any]] = {"axis": axis, "n", n}
     if prepend is not None:
         if prepend.device != x.device:
             raise ValueError(f"Arrays from two different devices ({prepend.device} and {x.device}) can not be combined.")