2
2
3
3
import operator
4
4
import warnings
5
- from collections .abc import Callable
6
- from typing import Any
5
+
6
+ # https://github.com/pylint-dev/pylint/issues/10112
7
+ from collections .abc import Callable # pylint: disable=import-error
8
+ from typing import ClassVar
7
9
8
10
from ._lib import _utils
9
11
from ._lib ._compat import (
12
14
is_dask_array ,
13
15
is_writeable_array ,
14
16
)
15
- from ._lib ._typing import Array , ModuleType
17
+ from ._lib ._typing import Array , Index , ModuleType , Untyped
16
18
17
19
__all__ = [
18
20
"at" ,
@@ -559,7 +561,7 @@ def sinc(x: Array, /, *, xp: ModuleType | None = None) -> Array:
559
561
_undef = object ()
560
562
561
563
562
- class at :
564
+ class at : # pylint: disable=invalid-name
563
565
"""
564
566
Update operations for read-only arrays.
565
567
@@ -651,14 +653,14 @@ class at:
651
653
"""
652
654
653
655
x : Array
654
- idx : Any
655
- __slots__ = ("idx" , "x" )
656
+ idx : Index
657
+ __slots__ : ClassVar [ tuple [ str , str ]] = ("idx" , "x" )
656
658
657
- def __init__ (self , x : Array , idx : Any = _undef , / ):
659
+ def __init__ (self , x : Array , idx : Index = _undef , / ):
658
660
self .x = x
659
661
self .idx = idx
660
662
661
- def __getitem__ (self , idx : Any ) -> Any :
663
+ def __getitem__ (self , idx : Index ) -> at :
662
664
"""Allow for the alternate syntax ``at(x)[start:stop:step]``,
663
665
which looks prettier than ``at(x, slice(start, stop, step))``
664
666
and feels more intuitive coming from the JAX documentation.
@@ -677,8 +679,8 @@ def _common(
677
679
copy : bool | None = True ,
678
680
xp : ModuleType | None = None ,
679
681
_is_update : bool = True ,
680
- ** kwargs : Any ,
681
- ) -> tuple [Any , None ] | tuple [None , Array ]:
682
+ ** kwargs : Untyped ,
683
+ ) -> tuple [Untyped , None ] | tuple [None , Array ]:
682
684
"""Perform common prepocessing.
683
685
684
686
Returns
@@ -706,11 +708,11 @@ def _common(
706
708
if not writeable :
707
709
msg = "Cannot modify parameter in place"
708
710
raise ValueError (msg )
709
- elif copy is None :
711
+ elif copy is None : # type: ignore[redundant-expr]
710
712
writeable = is_writeable_array (x )
711
713
copy = _is_update and not writeable
712
714
else :
713
- msg = f"Invalid value for copy: { copy !r} " # type: ignore[unreachable]
715
+ msg = f"Invalid value for copy: { copy !r} " # type: ignore[unreachable] # pyright: ignore[reportUnreachable]
714
716
raise ValueError (msg )
715
717
716
718
if copy :
@@ -741,7 +743,7 @@ def _common(
741
743
742
744
return None , x
743
745
744
- def get (self , ** kwargs : Any ) -> Any :
746
+ def get (self , ** kwargs : Untyped ) -> Untyped :
745
747
"""Return ``x[idx]``. In addition to plain ``__getitem__``, this allows ensuring
746
748
that the output is either a copy or a view; it also allows passing
747
749
keyword arguments to the backend.
@@ -766,7 +768,7 @@ def get(self, **kwargs: Any) -> Any:
766
768
assert x is not None
767
769
return x [self .idx ]
768
770
769
- def set (self , y : Array , / , ** kwargs : Any ) -> Array :
771
+ def set (self , y : Array , / , ** kwargs : Untyped ) -> Array :
770
772
"""Apply ``x[idx] = y`` and return the update array"""
771
773
res , x = self ._common ("set" , y , ** kwargs )
772
774
if res is not None :
@@ -781,7 +783,7 @@ def _iop(
781
783
elwise_op : Callable [[Array , Array ], Array ],
782
784
y : Array ,
783
785
/ ,
784
- ** kwargs : Any ,
786
+ ** kwargs : Untyped ,
785
787
) -> Array :
786
788
"""x[idx] += y or equivalent in-place operation on a subset of x
787
789
@@ -799,33 +801,33 @@ def _iop(
799
801
x [self .idx ] = elwise_op (x [self .idx ], y )
800
802
return x
801
803
802
- def add (self , y : Array , / , ** kwargs : Any ) -> Array :
804
+ def add (self , y : Array , / , ** kwargs : Untyped ) -> Array :
803
805
"""Apply ``x[idx] += y`` and return the updated array"""
804
806
return self ._iop ("add" , operator .add , y , ** kwargs )
805
807
806
- def subtract (self , y : Array , / , ** kwargs : Any ) -> Array :
808
+ def subtract (self , y : Array , / , ** kwargs : Untyped ) -> Array :
807
809
"""Apply ``x[idx] -= y`` and return the updated array"""
808
810
return self ._iop ("subtract" , operator .sub , y , ** kwargs )
809
811
810
- def multiply (self , y : Array , / , ** kwargs : Any ) -> Array :
812
+ def multiply (self , y : Array , / , ** kwargs : Untyped ) -> Array :
811
813
"""Apply ``x[idx] *= y`` and return the updated array"""
812
814
return self ._iop ("multiply" , operator .mul , y , ** kwargs )
813
815
814
- def divide (self , y : Array , / , ** kwargs : Any ) -> Array :
816
+ def divide (self , y : Array , / , ** kwargs : Untyped ) -> Array :
815
817
"""Apply ``x[idx] /= y`` and return the updated array"""
816
818
return self ._iop ("divide" , operator .truediv , y , ** kwargs )
817
819
818
- def power (self , y : Array , / , ** kwargs : Any ) -> Array :
820
+ def power (self , y : Array , / , ** kwargs : Untyped ) -> Array :
819
821
"""Apply ``x[idx] **= y`` and return the updated array"""
820
822
return self ._iop ("power" , operator .pow , y , ** kwargs )
821
823
822
- def min (self , y : Array , / , ** kwargs : Any ) -> Array :
824
+ def min (self , y : Array , / , ** kwargs : Untyped ) -> Array :
823
825
"""Apply ``x[idx] = minimum(x[idx], y)`` and return the updated array"""
824
826
xp = array_namespace (self .x )
825
827
y = xp .asarray (y )
826
828
return self ._iop ("min" , xp .minimum , y , ** kwargs )
827
829
828
- def max (self , y : Array , / , ** kwargs : Any ) -> Array :
830
+ def max (self , y : Array , / , ** kwargs : Untyped ) -> Array :
829
831
"""Apply ``x[idx] = maximum(x[idx], y)`` and return the updated array"""
830
832
xp = array_namespace (self .x )
831
833
y = xp .asarray (y )
0 commit comments