Add ruff to ci setup

.github/workflows/ruff.yml

@@ -0,0 +1,19 @@
+name: CI
+on: [push, pull_request]
+jobs:
+  check-ruff:
+    runs-on: ubuntu-latest
+    continue-on-error: true
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install ruff
+      # Update output format to enable automatic inline annotations.
+      - name: Run Ruff
+        run: ruff check --output-format=github --select F822,PLC0414,RUF022 --preview .

array_api_compat/__init__.py

@@ -19,4 +19,4 @@
 """
 __version__ = '1.4.1'
 
-from .common import *
+from .common import *  # noqa: F401, F403

array_api_compat/_internal.py

@@ -5,6 +5,7 @@
 from functools import wraps
 from inspect import signature
 
+
 def get_xp(xp):
     """
     Decorator to automatically replace xp with the corresponding array module.
@@ -21,13 +22,16 @@ def func(x, /, xp, kwarg=None):
     arguments.
 
     """
+
     def inner(f):
         @wraps(f)
         def wrapped_f(*args, **kwargs):
             return f(*args, xp=xp, **kwargs)
 
         sig = signature(f)
-        new_sig = sig.replace(parameters=[sig.parameters[i] for i in sig.parameters if i != 'xp'])
+        new_sig = sig.replace(
+            parameters=[sig.parameters[i] for i in sig.parameters if i != "xp"]
+        )
 
         if wrapped_f.__doc__ is None:
             wrapped_f.__doc__ = f"""\
@@ -41,3 +45,31 @@ def wrapped_f(*args, **kwargs):
         return wrapped_f
 
     return inner
+
+
+def _get_all_public_members(module, exclude=None, extend_all=False):
+    """Get all public members of a module.
+
+    Parameters
+    ----------
+    module : module
+        The module to get members from.
+    exclude : callable, optional
+        A callable that takes a name and returns True if the name should be
+        excluded from the list of members.
+    extend_all : bool, optional
+        If True, extend the module's __all__ attribute with the members of the
+        module derived from dir(module). To be used for libraries that do not have a complete __all__ list.
+    """
+    members = getattr(module, "__all__", [])
+
+    if members and not extend_all:
+        return members
+
+    if exclude is None:
+        exclude = lambda name: name.startswith("_")  # noqa: E731
+
+    members = members + [_ for _ in dir(module) if not exclude(_)]
+
+    # remove duplicates
+    return list(set(members))

array_api_compat/common/__init__.py

@@ -1 +1,17 @@
-from ._helpers import *
+from ._helpers import (
+    array_namespace,
+    device,
+    get_namespace,
+    is_array_api_obj,
+    size,
+    to_device,
+)
+
+__all__ = [
+    "array_namespace",
+    "device",
+    "get_namespace",
+    "is_array_api_obj",
+    "size",
+    "to_device",
+]

array_api_compat/common/_aliases.py

@@ -6,7 +6,8 @@
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from typing import Optional, Sequence, Tuple, Union, List
+    import numpy as np
+    from typing import Optional, Sequence, Tuple, Union
     from ._typing import ndarray, Device, Dtype, NestedSequence, SupportsBufferProtocol
 
 from typing import NamedTuple
@@ -544,11 +545,3 @@ def isdtype(
         # more strict here to match the type annotation? Note that the
         # numpy.array_api implementation will be very strict.
         return dtype == kind
-
-__all__ = ['arange', 'empty', 'empty_like', 'eye', 'full', 'full_like',
-           'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like',
-           'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
-           'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
-           'astype', 'std', 'var', 'permute_dims', 'reshape', 'argsort',
-           'sort', 'nonzero', 'sum', 'prod', 'ceil', 'floor', 'trunc',
-           'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype']

array_api_compat/common/_helpers.py

@@ -7,6 +7,12 @@
 """
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Optional, Union, Any
+    from ._typing import Array, Device  
+
 import sys
 import math
 
@@ -142,7 +148,7 @@ def _check_device(xp, device):
 # wrapping or subclassing them. These helper functions can be used instead of
 # the wrapper functions for libraries that need to support both NumPy/CuPy and
 # other libraries that use devices.
-def device(x: "Array", /) -> "Device":
+def device(x: Array, /) -> Device:
     """
     Hardware device the array data resides on.
 
@@ -204,7 +210,7 @@ def _torch_to_device(x, device, /, stream=None):
         raise NotImplementedError
     return x.to(device)
 
-def to_device(x: "Array", device: "Device", /, *, stream: "Optional[Union[int, Any]]" = None) -> "Array":
+def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]] = None) -> Array:
     """
     Copy the array from the device on which it currently resides to the specified ``device``.
 
@@ -252,5 +258,3 @@ def size(x):
     if None in x.shape:
         return None
     return math.prod(x.shape)
-
-__all__ = ['is_array_api_obj', 'array_namespace', 'get_namespace', 'device', 'to_device', 'size']

array_api_compat/common/_linalg.py

@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING, NamedTuple
 if TYPE_CHECKING:
-    from typing import Literal, Optional, Sequence, Tuple, Union
+    from typing import Literal, Optional, Tuple, Union
     from ._typing import ndarray
 
 import numpy as np
@@ -11,7 +11,7 @@
 else:
     from numpy.core.numeric import normalize_axis_tuple
 
-from ._aliases import matmul, matrix_transpose, tensordot, vecdot, isdtype
+from ._aliases import matrix_transpose, isdtype
 from .._internal import get_xp
 
 # These are in the main NumPy namespace but not in numpy.linalg
@@ -149,10 +149,4 @@ def trace(x: ndarray, /, xp, *, offset: int = 0, dtype=None, **kwargs) -> ndarra
             dtype = xp.float64
         elif x.dtype == xp.complex64:
             dtype = xp.complex128
-    return xp.asarray(xp.trace(x, offset=offset, dtype=dtype, axis1=-2, axis2=-1, **kwargs))
-
-__all__ = ['cross', 'matmul', 'outer', 'tensordot', 'EighResult',
-           'QRResult', 'SlogdetResult', 'SVDResult', 'eigh', 'qr', 'slogdet',
-           'svd', 'cholesky', 'matrix_rank', 'pinv', 'matrix_norm',
-           'matrix_transpose', 'svdvals', 'vecdot', 'vector_norm', 'diagonal',
-           'trace']
+    return xp.asarray(xp.trace(x, offset=offset, dtype=dtype, axis1=-2, axis2=-1, **kwargs))

array_api_compat/common/_typing.py

@@ -18,3 +18,6 @@ def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
     def __len__(self, /) -> int: ...
 
 SupportsBufferProtocol = Any
+
+Array = Any
+Device = Any

array_api_compat/cupy/__init__.py

@@ -1,16 +1,153 @@
-from cupy import *
+import cupy as _cp
+from cupy import *  # noqa: F401, F403
 
 # from cupy import * doesn't overwrite these builtin names
 from cupy import abs, max, min, round
 
+from .._internal import _get_all_public_members
+from ..common._helpers import (
+    array_namespace,
+    device,
+    get_namespace,
+    is_array_api_obj,
+    size,
+    to_device,
+)
+
 # These imports may overwrite names from the import * above.
-from ._aliases import *
+from ._aliases import (
+    UniqueAllResult,
+    UniqueCountsResult,
+    UniqueInverseResult,
+    acos,
+    acosh,
+    arange,
+    argsort,
+    asarray,
+    asarray_cupy,
+    asin,
+    asinh,
+    astype,
+    atan,
+    atan2,
+    atanh,
+    bitwise_invert,
+    bitwise_left_shift,
+    bitwise_right_shift,
+    bool,
+    ceil,
+    concat,
+    empty,
+    empty_like,
+    eye,
+    floor,
+    full,
+    full_like,
+    isdtype,
+    linspace,
+    matmul,
+    matrix_transpose,
+    nonzero,
+    ones,
+    ones_like,
+    permute_dims,
+    pow,
+    prod,
+    reshape,
+    sort,
+    std,
+    sum,
+    tensordot,
+    trunc,
+    unique_all,
+    unique_counts,
+    unique_inverse,
+    unique_values,
+    var,
+    vecdot,
+    zeros,
+    zeros_like,
+)
 
-# See the comment in the numpy __init__.py
-__import__(__package__ + '.linalg')
+__all__ = []
+
+__all__ += _get_all_public_members(_cp)
+
+__all__ += [
+    "abs",
+    "max",
+    "min",
+    "round",
+]
 
-from .linalg import matrix_transpose, vecdot
+__all__ += [
+    "array_namespace",
+    "device",
+    "get_namespace",
+    "is_array_api_obj",
+    "size",
+    "to_device",
+]
 
-from ..common._helpers import *
+__all__ += [
+    "UniqueAllResult",
+    "UniqueCountsResult",
+    "UniqueInverseResult",
+    "acos",
+    "acosh",
+    "arange",
+    "argsort",
+    "asarray",
+    "asarray_cupy",
+    "asin",
+    "asinh",
+    "astype",
+    "atan",
+    "atan2",
+    "atanh",
+    "bitwise_invert",
+    "bitwise_left_shift",
+    "bitwise_right_shift",
+    "bool",
+    "ceil",
+    "concat",
+    "empty",
+    "empty_like",
+    "eye",
+    "floor",
+    "full",
+    "full_like",
+    "isdtype",
+    "linspace",
+    "matmul",
+    "matrix_transpose",
+    "nonzero",
+    "ones",
+    "ones_like",
+    "permute_dims",
+    "pow",
+    "prod",
+    "reshape",
+    "sort",
+    "std",
+    "sum",
+    "tensordot",
+    "trunc",
+    "unique_all",
+    "unique_counts",
+    "unique_inverse",
+    "unique_values",
+    "var",
+    "zeros",
+    "zeros_like",
+]
+
+__all__ += [
+    "matrix_transpose",
+    "vecdot",
+]
+
+# See the comment in the numpy __init__.py
+__import__(__package__ + ".linalg")
 
-__array_api_version__ = '2022.12'
+__array_api_version__ = "2022.12"