Repository: scipy/scipy
File: scipy/sparse/_construct.py (Python)

Code before:
    axis1 : int
        First axis.
    axis2 : int
        Second axis.

    Returns
    -------
    a_swapped : sparse array in COO format
        A copy of the input array with the two identified axes swapped.

    Raises
    ------
    ValueError
        If provided a non-integer or out of range ``[-N, N-1]`` axis,
        where ``N`` is ``A.ndim``.
    """
    if not isintlike(axis1):
        raise ValueError(f"Invalid axis1 {axis1}. Must be an integer.")
    idx1 = axis1 if axis1 >= 0 else axis1 + A.ndim
    if idx1 < 0 or idx1 > A.ndim - 1:
        raise ValueError(f"Invalid axis1 {axis1} for N={A.ndim}. Must be in [-N, N-1].")
    if not isintlike(axis2):
        raise ValueError(f"Invalid axis2 {axis2}. Must be an integer.")
    idx2 = axis2 if axis2 >= 0 else axis2 + A.ndim
    if idx2 < 0 or idx2 > A.ndim - 1:
        raise ValueError(f"Invalid axis2 {axis2} for N={A.ndim}. Must be in [-N, N-1].")

    A = A.tocoo(copy=True)
    if idx1 == idx2:
        return A
    new_coords = list(A.coords)
    new_coords[idx1], new_coords[idx2] = new_coords[idx2], new_coords[idx1]
    A.coords = tuple(new_coords)
    new_shape = list(A.shape)
    new_shape[idx1], new_shape[idx2] = new_shape[idx2], new_shape[idx1]
    A._shape = tuple(new_shape)
    A.has_canonical_format = False  # data not guaranteed sorted
    return A


def permute_dims(A, axes=None, copy=False):
    """Permute the axes of the sparse array `A` to the order `axes`.

    Parameters
    ----------
    A : sparse array
    axes : tuple or list of ints, optional
        If specified, it must be a tuple or list which contains a permutation
        of ``[0, 1, ..., N-1]`` where ``N`` is the number of axes of `A`. The ``i``th
        axis of the returned array will correspond to the axis numbered ``axes[i]``
        of the input. If not specified, defaults to ``range(a.ndim)[::-1]``,

Review comment (refactor):
Side note: it would be nice if we had a helper function for this kind of integer axis validation. We have `_sputils.validateaxis` but that's more for the general case, and wouldn't produce very helpful error messages. Might be worth adding a related function like `validate_single_axis()` or something, that does the type checking, negative handling, bounds checking, and returns an int.

Code after:
    A : sparse array
    axis1 : int
        First axis.
    axis2 : int
        Second axis.

    Returns
    -------
    a_swapped : sparse array in COO format
        A copy of the input array with the two identified axes swapped.

    Raises
    ------
    ValueError
        If provided a non-integer or out of range ``[-N, N-1]`` axis,
        where ``N`` is ``A.ndim``.
    """
    axes = np.arange(A.ndim)
    try:
        axes[[axis1, axis2]] = axes[[axis2, axis1]]
    except IndexError as err:
        # original msg looks like: "index -4 is out of bounds for axis 0 with size 2"
        msg = str(err)
        msg = msg.removeprefix('index ').split(' axis ', 1)[0]
        # Final error is: "Invalid axis: -4 is out of bounds for ndim=2"
        raise ValueError(f"Invalid axis: {msg} ndim={A.ndim}")

    axes = tuple(axes)
    return permute_dims(A, axes=axes, copy=True)


def permute_dims(A, axes=None, copy=False):
    """Permute the axes of the sparse array `A` to the order `axes`.

    Parameters
    ----------
    A : sparse array
    axes : tuple or list of ints, optional
        If specified, it must be a tuple or list which contains a permutation
        of ``[0, 1, ..., N-1]`` where ``N`` is ``A.ndim``. The ith
        axis of the returned array will correspond to the axis numbered ``axes[i]``
        of the input. If not specified, defaults to ``range(A.ndim)[::-1]``,
        which reverses the order of the axes.
    copy : bool, optional (default: False)
        Whether to return the permutation as a copy. If False, an in-place
        permutation is provided if possible depending on format.

    Returns
    -------
    out : sparse array in COO format
        A copy of `A` with permuted axes.