ronantakizawa/codereview-bench huggingface dataset:
Repository: numpy/numpy
File: numpy/ma/core.pyi (Python)

Code before:
    *,
    keepdims: bool | _NoValueType = ...,
) -> _ArrayT: ...

minimum: _extrema_operation
maximum: _extrema_operation

@overload
def take(  # type: ignore[overload-overlap]
    a: _MaskedArray[_SCT],
    indices: _IntLike_co,
    axis: None = None,
    out: None = None,
    mode: _ModeKind = 'raise'
) -> _SCT: ...
@overload
def take(
    a: ArrayLike,
    indices: _IntLike_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> Any: ...
@overload
def take(
    a: _MaskedArray[_SCT],
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> _MaskedArray[_SCT]: ...
@overload
def take(
    a: ArrayLike,
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> _MaskedArray[Any]: ...
@overload
def take(
    a: ArrayLike,
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None,
    out: _ArrayT,
    mode: _ModeKind = 'raise',
) -> _ArrayT: ...
@overload
def take(
    a: ArrayLike,
    indices: _ArrayLikeInt_co,

Review comment (suggestion):
```pycon
>>> np.ma.take([[1,2], [3,4]], (1,2))
masked_array(data=[2, 3],
             mask=False,
       fill_value=999999)
```

so we can broaden `a`:

```suggestion
    a: _ArrayLike[_SCT],
```

Code after:
) -> _ArrayT: ...

minimum: _extrema_operation
maximum: _extrema_operation

@overload
def take(  # type: ignore[overload-overlap]
    a: _ArrayLike[_SCT],
    indices: _IntLike_co,
    axis: None = None,
    out: None = None,
    mode: _ModeKind = 'raise'
) -> _SCT: ...
@overload
def take(
    a: _ArrayLike[_SCT],
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> _MaskedArray[_SCT]: ...
@overload
def take(
    a: ArrayLike,
    indices: _IntLike_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> Any: ...
@overload
def take(
    a: ArrayLike,
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> _MaskedArray[Any]: ...
@overload
def take(
    a: ArrayLike,
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None,
    out: _ArrayT,
    mode: _ModeKind = 'raise',
) -> _ArrayT: ...
@overload
def take(
    a: ArrayLike,
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None = None,
    *,
Repository: numpy/numpy
File: numpy/ma/core.pyi (Python)

Code before:
    axis: SupportsIndex | None,
    fill_value: _ScalarLike_co | None,
    out: _ArrayT,
    *,
    keepdims: bool | _NoValueType = ...,
) -> _ArrayT: ...

minimum: _extrema_operation
maximum: _extrema_operation

@overload
def take(  # type: ignore[overload-overlap]
    a: _MaskedArray[_SCT],
    indices: _IntLike_co,
    axis: None = None,
    out: None = None,
    mode: _ModeKind = 'raise'
) -> _SCT: ...
@overload
def take(
    a: ArrayLike,
    indices: _IntLike_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> Any: ...
@overload
def take(
    a: _MaskedArray[_SCT],
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> _MaskedArray[_SCT]: ...
@overload
def take(
    a: ArrayLike,
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> _MaskedArray[Any]: ...
@overload
def take(
    a: ArrayLike,
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None,
    out: _ArrayT,
    mode: _ModeKind = 'raise',
) -> _ArrayT: ...
@overload

Review comment (suggestion):
This will match 

```pycon
>>> np.ma.take(np.eye(2), 1, axis=0)
masked_array(data=[0., 1.],
             mask=False,
       fill_value=1e+20)
```

and infer the return type as `Any`. 

By swapping this overload with the one below, then that overload would be matched instead, and we'd get a more precise return type of `_MaskedArray[_SCT]`. 

Code after:
    out: _ArrayT,
    *,
    keepdims: bool | _NoValueType = ...,
) -> _ArrayT: ...

minimum: _extrema_operation
maximum: _extrema_operation

@overload
def take(  # type: ignore[overload-overlap]
    a: _ArrayLike[_SCT],
    indices: _IntLike_co,
    axis: None = None,
    out: None = None,
    mode: _ModeKind = 'raise'
) -> _SCT: ...
@overload
def take(
    a: _ArrayLike[_SCT],
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> _MaskedArray[_SCT]: ...
@overload
def take(
    a: ArrayLike,
    indices: _IntLike_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> Any: ...
@overload
def take(
    a: ArrayLike,
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None = None,
    out: None = None,
    mode: _ModeKind = 'raise',
) -> _MaskedArray[Any]: ...
@overload
def take(
    a: ArrayLike,
    indices: _ArrayLikeInt_co,
    axis: SupportsIndex | None,
    out: _ArrayT,
    mode: _ModeKind = 'raise',
) -> _ArrayT: ...
@overload
def take(
    a: ArrayLike,
Repository: numpy/numpy
File: numpy/ma/core.pyi (Python)

Code before:
    def ptp(
        self,
        axis: _ShapeLike | None = None,
        *,
        out: _ArrayT,
        fill_value: _ScalarLike_co | None = None,
        keepdims: bool = False,
    ) -> _ArrayT: ...

    # Keep in-sync with np.ma.take
    @overload
    def take(  # type: ignore[overload-overlap]
        self: _MaskedArray[_SCT],
        indices: _IntLike_co,
        axis: None = None,
        out: None = None,
        mode: _ModeKind = 'raise'
    ) -> _SCT: ...
    @overload
    def take(
        self,
        indices: _IntLike_co,
        axis: SupportsIndex | None = None,
        out: None = None,
        mode: _ModeKind = 'raise',
    ) -> Any: ...
    @overload
    def take(
        self: _MaskedArray[_SCT],
        indices: _ArrayLikeInt_co,
        axis: SupportsIndex | None = None,
        out: None = None,
        mode: _ModeKind = 'raise',
    ) -> _MaskedArray[_SCT]: ...
    @overload
    def take(
        self,
        indices: _ArrayLikeInt_co,
        axis: SupportsIndex | None = None,
        out: None = None,
        mode: _ModeKind = 'raise',
    ) -> _MaskedArray[Any]: ...
    @overload
    def take(
        self,
        indices: _ArrayLikeInt_co,
        axis: SupportsIndex | None,
        out: _ArrayT,
        mode: _ModeKind = 'raise',
    ) -> _ArrayT: ...
    @overload

Review comment (suggestion):
The overload above always gets matched, so unless I'm missing something, this one can be removed.

```suggestion
```

Code after:
    def ptp(
        self,
        axis: _ShapeLike | None = None,
        *,
        out: _ArrayT,
        fill_value: _ScalarLike_co | None = None,
        keepdims: bool = False,
    ) -> _ArrayT: ...

    #
    def partition(
        self,
        kth: _ArrayLikeInt,
        axis: SupportsIndex = -1,
        kind: _PartitionKind = "introselect",
        order: str | Sequence[str] | None = None
    ) -> None: ...
    def argpartition(
        self,
        kth: _ArrayLikeInt,
        axis: SupportsIndex = -1,
        kind: _PartitionKind = "introselect",
        order: str | Sequence[str] | None = None
    ) -> _MaskedArray[intp]: ...

    # Keep in-sync with np.ma.take
    @overload
    def take(  # type: ignore[overload-overlap]
        self: _MaskedArray[_SCT],
        indices: _IntLike_co,
        axis: None = None,
        out: None = None,
        mode: _ModeKind = 'raise'
    ) -> _SCT: ...
    @overload
    def take(
        self: _MaskedArray[_SCT],
        indices: _ArrayLikeInt_co,
        axis: SupportsIndex | None = None,
        out: None = None,
        mode: _ModeKind = 'raise',
    ) -> _MaskedArray[_SCT]: ...
    @overload
    def take(
        self,
        indices: _ArrayLikeInt_co,
        axis: SupportsIndex | None,
        out: _ArrayT,
        mode: _ModeKind = 'raise',
    ) -> _ArrayT: ...
    @overload
Repository: numpy/numpy
File: numpy/_core/src/multiarray/temp_elide.c (C)

Code before:
#endif

/*
 * linear search pointer in table
 * number of pointers is usually quite small but if a performance impact can be
 * measured this could be converted to a binary search
 */
static int
find_addr(void * addresses[], npy_intp naddr, void * addr)
{
    npy_intp j;
    for (j = 0; j < naddr; j++) {
        if (addr == addresses[j]) {
            return 1;
        }
    }
    return 0;
}

static int
check_unique_temporary(PyObject *lhs)
{
#if PY_VERSION_HEX >= 0x030E00A7
    // In Python 3.14.0a7 and later, a reference count of one doesn't guarantee
    // that an object is a unique temporary variable. We scan the top of the
    // interpreter stack to check if the object exists as an "onwned" reference
    // in the temporary stack.
    PyFrameObject *frame = PyEval_GetFrame();
    if (frame == NULL) {
        return 0;
    }
    _PyInterpreterFrame *f = frame->f_frame;
    _PyStackRef *base = _PyFrame_Stackbase(f);
    _PyStackRef *stackpointer = f->stackpointer;
    int found_once = 0;
    while (stackpointer > base) {
        stackpointer--;
        PyObject *obj = PyStackRef_AsPyObjectBorrow(*stackpointer);
        if (obj == lhs) {
            if (!found_once && PyStackRef_IsHeapSafe(*stackpointer)) {
                found_once = 1;
            }
            else {
                return 0;
            }
        }
        return found_once;
    }
    return 0;
#else
    return 1;

Review comment (suggestion):
```suggestion
    // interpreter stack to check if the object exists as an "owned" reference
```

Code after:
#endif

/*
 * linear search pointer in table
 * number of pointers is usually quite small but if a performance impact can be
 * measured this could be converted to a binary search
 */
static int
find_addr(void * addresses[], npy_intp naddr, void * addr)
{
    npy_intp j;
    for (j = 0; j < naddr; j++) {
        if (addr == addresses[j]) {
            return 1;
        }
    }
    return 0;
}

static int
check_unique_temporary(PyObject *lhs)
{
#if PY_VERSION_HEX >= 0x030E00A7 && !defined(PYPY_VERSION)
    // In Python 3.14.0a7 and later, a reference count of one doesn't guarantee
    // that an object is a unique temporary variable. We scan the top of the
    // interpreter stack to check if the object exists as an "owned" reference
    // in the temporary stack.
    PyFrameObject *frame = PyEval_GetFrame();
    if (frame == NULL) {
        return 0;
    }
    _PyInterpreterFrame *f = frame->f_frame;
    _PyStackRef *base = _PyFrame_Stackbase(f);
    _PyStackRef *stackpointer = f->stackpointer;
    int found_once = 0;
    while (stackpointer > base) {
        stackpointer--;
        PyObject *obj = PyStackRef_AsPyObjectBorrow(*stackpointer);
        if (obj == lhs) {
            if (!found_once && PyStackRef_IsHeapSafe(*stackpointer)) {
                found_once = 1;
            }
            else {
                return 0;
            }
        }
        return found_once;
    }
    return 0;
#else
    return 1;
Repository: numpy/numpy
File: numpy/_core/src/multiarray/temp_elide.c (C)

Code before:
/*
 * linear search pointer in table
 * number of pointers is usually quite small but if a performance impact can be
 * measured this could be converted to a binary search
 */
static int
find_addr(void * addresses[], npy_intp naddr, void * addr)
{
    npy_intp j;
    for (j = 0; j < naddr; j++) {
        if (addr == addresses[j]) {
            return 1;
        }
    }
    return 0;
}

static int
check_unique_temporary(PyObject *lhs)
{
#if PY_VERSION_HEX >= 0x030E00A7
    // In Python 3.14.0a7 and later, a reference count of one doesn't guarantee
    // that an object is a unique temporary variable. We scan the top of the
    // interpreter stack to check if the object exists as an "owned" reference
    // in the temporary stack.
    PyFrameObject *frame = PyEval_GetFrame();
    if (frame == NULL) {
        return 0;
    }
    _PyInterpreterFrame *f = frame->f_frame;
    _PyStackRef *base = _PyFrame_Stackbase(f);
    _PyStackRef *stackpointer = f->stackpointer;
    int found_once = 0;
    while (stackpointer > base) {
        stackpointer--;
        PyObject *obj = PyStackRef_AsPyObjectBorrow(*stackpointer);
        if (obj == lhs) {
            if (!found_once && PyStackRef_IsHeapSafe(*stackpointer)) {
                found_once = 1;
            }
            else {
                return 0;
            }
        }
        return found_once;
    }
    return 0;
#else
    return 1;
#endif
}

Review comment (suggestion):
You can use `_PyEval_GetFrame` to get `_PyInterpreterFrame` directly here of the current frame without having to materialize a Python level frame object, this would also be more performant.

Code after:
/*
 * linear search pointer in table
 * number of pointers is usually quite small but if a performance impact can be
 * measured this could be converted to a binary search
 */
static int
find_addr(void * addresses[], npy_intp naddr, void * addr)
{
    npy_intp j;
    for (j = 0; j < naddr; j++) {
        if (addr == addresses[j]) {
            return 1;
        }
    }
    return 0;
}

static int
check_unique_temporary(PyObject *lhs)
{
#if PY_VERSION_HEX >= 0x030E00A7 && !defined(PYPY_VERSION)
    // In Python 3.14.0a7 and later, a reference count of one doesn't guarantee
    // that an object is a unique temporary variable. We scan the top of the
    // interpreter stack to check if the object exists as an "owned" reference
    // in the temporary stack.
    PyFrameObject *frame = PyEval_GetFrame();
    if (frame == NULL) {
        return 0;
    }
    _PyInterpreterFrame *f = frame->f_frame;
    _PyStackRef *base = _PyFrame_Stackbase(f);
    _PyStackRef *stackpointer = f->stackpointer;
    int found_once = 0;
    while (stackpointer > base) {
        stackpointer--;
        PyObject *obj = PyStackRef_AsPyObjectBorrow(*stackpointer);
        if (obj == lhs) {
            if (!found_once && PyStackRef_IsHeapSafe(*stackpointer)) {
                found_once = 1;
            }
            else {
                return 0;
            }
        }
        return found_once;
    }
    return 0;
#else
    return 1;
#endif
}
Repository: numpy/numpy
File: numpy/_core/tests/test_array_coercion.py (Python)

Code before:
    If full is True (default) it includes array-likes not capable of handling
    all dtypes.
    """
    # base array:
    def ndarray(a):
        return a

    yield param(ndarray, id="ndarray")

    # subclass:
    class MyArr(np.ndarray):
        pass

    def subclass(a):
        return a.view(MyArr)

    yield subclass

    class _SequenceLike:
        # Older NumPy versions, sometimes cared whether a protocol array was
        # also _SequenceLike.  This shouldn't matter, but keep it for now
        # for __array__ and not the others.
        def __len__(self):
            raise TypeError

        def __getitem__(self):  # noqa: PLE0302
            raise TypeError

    # Array-interface
    class ArrayDunder(_SequenceLike):
        def __init__(self, a):
            self.a = a

        def __array__(self, dtype=None, copy=None):
            if dtype is None:
                return self.a
            return self.a.astype(dtype)

    yield param(ArrayDunder, id="__array__")

    # memory-view
    yield param(memoryview, id="memoryview")

    # Array-interface
    class ArrayInterface:
        def __init__(self, a):
            self.a = a  # need to hold on to keep interface valid
            self.__array_interface__ = a.__array_interface__

    yield param(ArrayInterface, id="__array_interface__")


Review comment (suggestion):
might as well do

```suggestion
        def __getitem__(self, _, /):
```

because it'll `raise TypeError` either way

Code after:
    If full is True (default) it includes array-likes not capable of handling
    all dtypes.
    """
    # base array:
    def ndarray(a):
        return a

    yield param(ndarray, id="ndarray")

    # subclass:
    class MyArr(np.ndarray):
        pass

    def subclass(a):
        return a.view(MyArr)

    yield subclass

    class _SequenceLike:
        # Older NumPy versions, sometimes cared whether a protocol array was
        # also _SequenceLike.  This shouldn't matter, but keep it for now
        # for __array__ and not the others.
        def __len__(self):
            raise TypeError

        def __getitem__(self, _, /):
            raise TypeError

    # Array-interface
    class ArrayDunder(_SequenceLike):
        def __init__(self, a):
            self.a = a

        def __array__(self, dtype=None, copy=None):
            if dtype is None:
                return self.a
            return self.a.astype(dtype)

    yield param(ArrayDunder, id="__array__")

    # memory-view
    yield param(memoryview, id="memoryview")

    # Array-interface
    class ArrayInterface:
        def __init__(self, a):
            self.a = a  # need to hold on to keep interface valid
            self.__array_interface__ = a.__array_interface__

    yield param(ArrayInterface, id="__array_interface__")
Repository: numpy/numpy
File: numpy/_core/fromnumeric.pyi (Python)

Code before:

@overload
def transpose(
    a: _ArrayLike[_SCT],
    axes: _ShapeLike | None = ...
) -> NDArray[_SCT]: ...
@overload
def transpose(
    a: ArrayLike,
    axes: _ShapeLike | None = ...
) -> NDArray[Any]: ...

@overload
def matrix_transpose(x: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
@overload
def matrix_transpose(x: ArrayLike, /) -> NDArray[Any]: ...

#
@overload
def partition(
    a: _ArrayLike[_ScalarT],
    kth: _ArrayLikeInt,
    axis: SupportsIndex | None = -1,
    kind: _PartitionKind = "introselect",
    order: None = None,
) -> NDArray[_ScalarT]: ...
@overload
def partition(
    a: _ArrayLike[np.void],
    kth: _ArrayLikeInt,
    axis: SupportsIndex | None = -1,
    kind: _PartitionKind = "introselect",
    order: str | Sequence[str] | None = None,
) -> NDArray[np.void]: ...
@overload
def partition(
    a: ArrayLike,
    kth: _ArrayLikeInt,
    axis: SupportsIndex | None = -1,
    kind: _PartitionKind = "introselect",
    order: str | Sequence[str] | None = None,
) -> NDArray[Any]: ...

#
def argpartition(
    a: ArrayLike,
    kth: _ArrayLikeInt,
    axis: SupportsIndex | None = -1,
    kind: _PartitionKind = "introselect",
    order: str | Sequence[str] | None = None,
) -> NDArray[intp]: ...

Review comment (suggestion):
This typevar has been renamed in the meantime. This should fix the mypy errors

```suggestion
    a: _ArrayLike[_SCT],
    kth: _ArrayLikeInt,
    axis: SupportsIndex | None = -1,
    kind: _PartitionKind = "introselect",
    order: None = None,
) -> NDArray[_SCT]: ...
```

Code after:
) -> NDArray[Any]: ...

@overload
def transpose(
    a: _ArrayLike[_SCT],
    axes: _ShapeLike | None = ...
) -> NDArray[_SCT]: ...
@overload
def transpose(
    a: ArrayLike,
    axes: _ShapeLike | None = ...
) -> NDArray[Any]: ...

@overload
def matrix_transpose(x: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
@overload
def matrix_transpose(x: ArrayLike, /) -> NDArray[Any]: ...

#
@overload
def partition(
    a: _ArrayLike[_SCT],
    kth: _ArrayLikeInt,
    axis: SupportsIndex | None = -1,
    kind: _PartitionKind = "introselect",
    order: None = None,
) -> NDArray[_SCT]: ...
@overload
def partition(
    a: _ArrayLike[np.void],
    kth: _ArrayLikeInt,
    axis: SupportsIndex | None = -1,
    kind: _PartitionKind = "introselect",
    order: str | Sequence[str] | None = None,
) -> NDArray[np.void]: ...
@overload
def partition(
    a: ArrayLike,
    kth: _ArrayLikeInt,
    axis: SupportsIndex | None = -1,
    kind: _PartitionKind = "introselect",
    order: str | Sequence[str] | None = None,
) -> NDArray[Any]: ...

#
def argpartition(
    a: ArrayLike,
    kth: _ArrayLikeInt,
    axis: SupportsIndex | None = -1,
    kind: _PartitionKind = "introselect",
    order: str | Sequence[str] | None = None,
Repository: numpy/numpy
File: numpy/_core/numeric.pyi (Python)

Code before:

    # other
    _Array,
    _ConstructorEmpty,
    _KwargsEmpty,
)

from numpy._typing import (
    ArrayLike,
    NDArray,
    DTypeLike,
    _SupportsDType,
    _ShapeLike,
    _DTypeLike,
    _ArrayLike,
    _SupportsArrayFunc,
    _ScalarLike_co,
    _ArrayLikeBool_co,
    _ArrayLikeUInt_co,
    _ArrayLikeInt_co,
    _ArrayLikeFloat_co,
    _ArrayLikeComplex_co,
    _ArrayLikeTD64_co,
    _ArrayLikeObject_co,
    _NestedSequence,
)

__all__ = [
    "newaxis",
    "ndarray",
    "flatiter",
    "nditer",
    "nested_iters",
    "ufunc",
    "arange",
    "array",
    "asarray",
    "asanyarray",
    "ascontiguousarray",
    "asfortranarray",
    "zeros",
    "count_nonzero",
    "empty",
    "broadcast",
    "dtype",
    "fromstring",
    "fromfile",
    "frombuffer",
    "from_dlpack",
    "where",
    "argwhere",

Review comment (suggestion):
```suggestion
    _ArrayLikeUnknown,
    _NestedSequence,
```

Code after:

    # other
    _Array,
    _ConstructorEmpty,
    _KwargsEmpty,
)

from numpy._typing import (
    ArrayLike,
    NDArray,
    DTypeLike,
    _SupportsDType,
    _ShapeLike,
    _DTypeLike,
    _ArrayLike,
    _SupportsArrayFunc,
    _ScalarLike_co,
    _ArrayLikeBool_co,
    _ArrayLikeUInt_co,
    _ArrayLikeInt_co,
    _ArrayLikeFloat_co,
    _ArrayLikeComplex_co,
    _ArrayLikeTD64_co,
    _ArrayLikeObject_co,
    _ArrayLikeUnknown,
    _NestedSequence,
)

__all__ = [
    "newaxis",
    "ndarray",
    "flatiter",
    "nditer",
    "nested_iters",
    "ufunc",
    "arange",
    "array",
    "asarray",
    "asanyarray",
    "ascontiguousarray",
    "asfortranarray",
    "zeros",
    "count_nonzero",
    "empty",
    "broadcast",
    "dtype",
    "fromstring",
    "fromfile",
    "frombuffer",
    "from_dlpack",
    "where",