Repository: numpy/numpy
File: numpy/lib/_arraypad_impl.py (Python)

Code before:
    >>> np.pad(a, {1: (1, 2)})
    array([[0, 0, 1, 2, 0, 0],
           [0, 3, 4, 5, 0, 0]])
    >>> np.pad(a, {-1: 2})
    array([[0, 0, 0, 1, 2, 0, 0],
           [0, 0, 3, 4, 5, 0, 0]])
    >>> np.pad(a, {0: (3, 0)})
    array([[0, 0, 0],
           [0, 0, 0],
           [0, 0, 0],
           [0, 1, 2],
           [3, 4, 5]])
    >>> np.pad(a, {0: (3, 0), 1: 2})
    array([[0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 1, 2, 0, 0],
           [0, 0, 3, 4, 5, 0, 0]])
    """
    if isinstance(pad_width, dict):
        seq = [(0, 0)] * np.ndim(array)
        for axis, width in pad_width.items():
            if isinstance(width, tuple):
                seq[axis] = width
            else:
                seq[axis] = (width, width)
        pad_width = seq
    array = np.asarray(array)
    pad_width = np.asarray(pad_width)

    if not pad_width.dtype.kind == 'i':
        raise TypeError('`pad_width` must be of integral type.')

    # Broadcast to shape (array.ndim, 2)
    pad_width = _as_pairs(pad_width, array.ndim, as_index=True)

    if callable(mode):
        # Old behavior: Use user-supplied function with np.apply_along_axis
        function = mode
        # Create a new zero padded array
        padded, _ = _pad_simple(array, pad_width, fill_value=0)
        # And apply along each axis

        for axis in range(padded.ndim):
            # Iterate using ndindex as in apply_along_axis, but assuming that
            # function operates inplace on the padded array.

            # view with the iteration axis at the end
            view = np.moveaxis(padded, axis, -1)

            # compute indices for the iteration axes, and append a trailing

Review comment (suggestion):
something like this would be more type-safe (and also a bit of an excuse to be able to use structural pattern matching for once haha)

```suggestion
            match width:
                case int(both):
                    seq[axis] = both, both
                case tuple((int(before), int(after))):
                    seq[axis] = before, after
                case _ as invalid:
                    typing.assert_never(invalid)
```

Code after:
           [100, 100, 100, 100, 100, 100, 100]])

    >>> a = np.arange(1, 7).reshape(2, 3)
    >>> np.pad(a, {1: (1, 2)})
    array([[0, 1, 2, 3, 0, 0],
           [0, 4, 5, 6, 0, 0]])
    >>> np.pad(a, {-1: 2})
    array([[0, 0, 1, 2, 3, 0, 0],
           [0, 0, 4, 5, 6, 0, 0]])
    >>> np.pad(a, {0: (3, 0)})
    array([[0, 0, 0],
           [0, 0, 0],
           [0, 0, 0],
           [1, 2, 3],
           [4, 5, 6]])
    >>> np.pad(a, {0: (3, 0), 1: 2})
    array([[0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0, 0, 0],
           [0, 0, 1, 2, 3, 0, 0],
           [0, 0, 4, 5, 6, 0, 0]])
    """
    array = np.asarray(array)
    if isinstance(pad_width, dict):
        seq = [(0, 0)] * array.ndim
        for axis, width in pad_width.items():
            match width:
                case int(both):
                    seq[axis] = both, both
                case tuple((int(before), int(after))):
                    seq[axis] = before, after
                case _ as invalid:
                    typing.assert_never(invalid)
        pad_width = seq
    pad_width = np.asarray(pad_width)

    if not pad_width.dtype.kind == 'i':
        raise TypeError('`pad_width` must be of integral type.')

    # Broadcast to shape (array.ndim, 2)
    pad_width = _as_pairs(pad_width, array.ndim, as_index=True)

    if callable(mode):
        # Old behavior: Use user-supplied function with np.apply_along_axis
        function = mode
        # Create a new zero padded array
        padded, _ = _pad_simple(array, pad_width, fill_value=0)
        # And apply along each axis

        for axis in range(padded.ndim):
            # Iterate using ndindex as in apply_along_axis, but assuming that