Add a test for __all__ self-consistency

Author: asmeurer

array_api_compat/common/_aliases.py

@@ -146,6 +146,9 @@ def zeros_like(
 
 # The functions here return namedtuples (np.unique() returns a normal
 # tuple).
+
+# Note that these named tuples aren't actually part of the standard namespace,
+# but I don't see any issue with exporting the names here regardless.
 class UniqueAllResult(NamedTuple):
     values: ndarray
     indices: ndarray

array_api_compat/common/_helpers.py

@@ -302,3 +302,5 @@ def size(x):
     "size",
     "to_device",
 ]
+
+_all_ignore = ['sys', 'math', 'inspect']

array_api_compat/dask/array/_aliases.py

@@ -50,7 +50,7 @@
 # an error with dask
 
 # TODO: delete the xp stuff, it shouldn't be necessary
-def dask_arange(
+def _dask_arange(
     start: Union[int, float],
     /,
     stop: Optional[Union[int, float]] = None,
@@ -72,7 +72,7 @@ def dask_arange(
     args.append(step)
     return xp.arange(*args, dtype=dtype, **kwargs)
 
-arange = get_xp(da)(dask_arange)
+arange = get_xp(da)(_dask_arange)
 eye = get_xp(da)(_aliases.eye)
 
 from functools import partial
@@ -142,4 +142,4 @@ def dask_arange(
                             'int16', 'int32', 'int64', 'uint8', 'uint16', 'uint32', 'uint64',
                             'complex64', 'complex128', 'iinfo', 'finfo', 'can_cast', 'result_type']
 
-del da, partial, common_aliases, _da_unsupported,
+_all_ignore = ['get_xp', 'da', 'partial', 'common_aliases', 'np']

array_api_compat/dask/array/linalg.py

@@ -51,6 +51,4 @@ def svdvals(x: ndarray) -> Union[ndarray, Tuple[ndarray, ...]]:
                         "cholesky", "matrix_rank", "matrix_norm", "svdvals",
                         "vector_norm", "diagonal"]
 
-del get_xp
-del da
-del _linalg
+_all_ignore = ['get_xp', 'da', 'linalg_all']

array_api_compat/numpy/_aliases.py

@@ -77,3 +77,5 @@
                               'acosh', 'asin', 'asinh', 'atan', 'atan2',
                               'atanh', 'bitwise_left_shift', 'bitwise_invert',
                               'bitwise_right_shift', 'concat', 'pow']
+
+_all_ignore = ['np', 'get_xp']

array_api_compat/torch/_aliases.py

@@ -1,11 +1,9 @@
 from __future__ import annotations
 
-from functools import wraps
-from builtins import all as builtin_all, any as builtin_any
+from functools import wraps as _wraps
+from builtins import all as _builtin_all, any as _builtin_any
 
-from ..common._aliases import (UniqueAllResult, UniqueCountsResult,
-                               UniqueInverseResult,
-                               matrix_transpose as _aliases_matrix_transpose,
+from ..common._aliases import (matrix_transpose as _aliases_matrix_transpose,
                                vecdot as _aliases_vecdot)
 from .._internal import get_xp
 
@@ -86,7 +84,7 @@
 
 
 def _two_arg(f):
-    @wraps(f)
+    @_wraps(f)
     def _f(x1, x2, /, **kwargs):
         x1, x2 = _fix_promotion(x1, x2)
         return f(x1, x2, **kwargs)
@@ -509,7 +507,7 @@ def arange(start: Union[int, float],
         start, stop = 0, start
     if step > 0 and stop <= start or step < 0 and stop >= start:
         if dtype is None:
-            if builtin_all(isinstance(i, int) for i in [start, stop, step]):
+            if _builtin_all(isinstance(i, int) for i in [start, stop, step]):
                 dtype = torch.int64
             else:
                 dtype = torch.float32
@@ -601,6 +599,11 @@ def broadcast_arrays(*arrays: array) -> List[array]:
     shape = torch.broadcast_shapes(*[a.shape for a in arrays])
     return [torch.broadcast_to(a, shape) for a in arrays]
 
+# Note that these named tuples aren't actually part of the standard namespace,
+# but I don't see any issue with exporting the names here regardless.
+from ..common._aliases import (UniqueAllResult, UniqueCountsResult,
+                               UniqueInverseResult)
+
 # https://github.com/pytorch/pytorch/issues/70920
 def unique_all(x: array) -> UniqueAllResult:
     # torch.unique doesn't support returning indices.
@@ -665,7 +668,7 @@ def isdtype(
     for more details
     """
     if isinstance(kind, tuple) and _tuple:
-        return builtin_any(isdtype(dtype, k, _tuple=False) for k in kind)
+        return _builtin_any(isdtype(dtype, k, _tuple=False) for k in kind)
     elif isinstance(kind, str):
         if kind == 'bool':
             return dtype == torch.bool
@@ -693,15 +696,19 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None, **kwargs) -
         axis = 0
     return torch.index_select(x, axis, indices, **kwargs)
 
-__all__ = ['result_type', 'can_cast', 'permute_dims', 'bitwise_invert', 'newaxis',
-           'add', 'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or',
-           'bitwise_right_shift', 'bitwise_xor', 'divide', 'equal',
-           'floor_divide', 'greater', 'greater_equal', 'less', 'less_equal',
-           'logaddexp', 'multiply', 'not_equal', 'pow', 'remainder',
-           'subtract', 'max', 'min', 'sort', 'prod', 'sum', 'any', 'all',
-           'mean', 'std', 'var', 'concat', 'squeeze', 'broadcast_to', 'flip', 'roll',
-           'nonzero', 'where', 'reshape', 'arange', 'eye', 'linspace', 'full',
-           'ones', 'zeros', 'empty', 'tril', 'triu', 'expand_dims', 'astype',
-           'broadcast_arrays', 'unique_all', 'unique_counts',
-           'unique_inverse', 'unique_values', 'matmul', 'matrix_transpose',
-           'vecdot', 'tensordot', 'isdtype', 'take']
+__all__ = ['result_type', 'can_cast', 'permute_dims', 'bitwise_invert',
+           'newaxis', 'add', 'atan2', 'bitwise_and', 'bitwise_left_shift',
+           'bitwise_or', 'bitwise_right_shift', 'bitwise_xor', 'divide',
+           'equal', 'floor_divide', 'greater', 'greater_equal', 'less',
+           'less_equal', 'logaddexp', 'multiply', 'not_equal', 'pow',
+           'remainder', 'subtract', 'max', 'min', 'sort', 'prod', 'sum',
+           'any', 'all', 'mean', 'std', 'var', 'concat', 'squeeze',
+           'broadcast_to', 'flip', 'roll', 'nonzero', 'where', 'reshape',
+           'arange', 'eye', 'linspace', 'full', 'ones', 'zeros', 'empty',
+           'tril', 'triu', 'expand_dims', 'astype', 'broadcast_arrays',
+           'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
+           'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
+           'matmul', 'matrix_transpose', 'vecdot', 'tensordot', 'isdtype',
+           'take']
+
+_all_ignore = ['torch', 'get_xp']

array_api_compat/torch/linalg.py

@@ -7,6 +7,8 @@
     from torch import dtype as Dtype
     from typing import Optional
 
+from ._aliases import _fix_promotion, sum
+
 from torch.linalg import * # noqa: F403
 
 # torch.linalg doesn't define __all__
@@ -16,7 +18,7 @@
 
 # outer is implemented in torch but aren't in the linalg namespace
 from torch import outer
-from ._aliases import _fix_promotion, matrix_transpose, tensordot, sum
+from ._aliases import matrix_transpose, tensordot
 
 # Note: torch.linalg.cross does not default to axis=-1 (it defaults to the
 # first axis with size 3), see https://github.com/pytorch/pytorch/issues/58743
@@ -59,4 +61,6 @@ def trace(x: array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> arr
 __all__ = linalg_all + ['outer', 'trace', 'matrix_transpose', 'tensordot',
                         'vecdot', 'solve']
 
+_all_ignore = ['torch_linalg', 'sum']
+
 del linalg_all

tests/test_all.py

@@ -0,0 +1,42 @@
+"""
+Test that files that define __all__ aren't missing any exports.
+
+You can add names that shouldn't be exported to _all_ignore, like
+
+_all_ignore = ['sys']
+
+This is preferable to del-ing the names as this will break any name that is
+used inside of a function. Note that names starting with an underscore are automatically ignored.
+"""
+
+
+import sys
+
+from ._helpers import import_
+
+import pytest
+
+@pytest.mark.parametrize("library", ["common", "cupy", "numpy", "torch", "dask.array"])
+def test_all(library):
+    import_(library, wrapper=True)
+
+    for mod_name in sys.modules:
+        if 'array_api_compat.' + library not in mod_name:
+            continue
+
+        module = sys.modules[mod_name]
+
+        # TODO: We should define __all__ in the __init__.py files and test it
+        # there too.
+        if not hasattr(module, '__all__'):
+            continue
+
+        dir_names = [n for n in dir(module) if not n.startswith('_')]
+        ignore_all_names = getattr(module, '_all_ignore', [])
+        ignore_all_names += ['annotations', 'TYPE_CHECKING']
+        dir_names = set(dir_names) - set(ignore_all_names)
+        all_names = module.__all__
+
+        if set(dir_names) != set(all_names):
+            assert set(dir_names) - set(all_names) == set(), f"Some dir() names not included in __all__ for {mod_name}"
+            assert set(all_names) - set(dir_names) == set(), f"Some __all__ names not in dir() for {mod_name}"