TYP: __dir__

Author: jorenham

array_api_compat/_internal.py

@@ -8,8 +8,6 @@
 from types import ModuleType
 from typing import TypeVar
 
-__all__ = ["get_xp"]
-
 _T = TypeVar("_T")
 
 
@@ -52,3 +50,10 @@ def wrapped_f(*args: object, **kwargs: object) -> object:
         return wrapped_f  # pyright: ignore[reportReturnType]
 
     return inner
+
+
+__all__ = ["get_xp"]
+
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/common/__init__.py

@@ -1 +1 @@
-from ._helpers import * # noqa: F403
+from ._helpers import *  # noqa: F403

array_api_compat/common/_aliases.py

@@ -720,5 +720,8 @@ def iinfo(type_: DType | Array, /, xp: Namespace) -> Any:
     "finfo",
     "iinfo",
 ]
-
 _all_ignore = ["inspect", "array_namespace", "NamedTuple"]
+
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/common/_fft.py

@@ -208,3 +208,6 @@ def ifftshift(
     "fftshift",
     "ifftshift",
 ]
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/common/_helpers.py

@@ -1044,3 +1044,6 @@ def is_lazy_array(x: object) -> bool:
 ]
 
 _all_ignore = ["sys", "math", "inspect", "warnings"]
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/common/_linalg.py

@@ -226,3 +226,7 @@ def trace(
            'trace']
 
 _all_ignore = ['math', 'normalize_axis_tuple', 'get_xp', 'np', 'isdtype']
+
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/common/_typing.py

@@ -3,17 +3,6 @@
 from types import ModuleType as Namespace
 from typing import Any, Protocol, TypeAlias, TypeVar
 
-__all__ = [
-    "Array",
-    "SupportsArrayNamespace",
-    "DType",
-    "Device",
-    "HasShape",
-    "Namespace",
-    "NestedSequence",
-    "SupportsBufferProtocol",
-]
-
 _T_co = TypeVar("_T_co", covariant=True)
 
 class NestedSequence(Protocol[_T_co]):
@@ -34,3 +23,19 @@ def shape(self, /) -> _T_co: ...
 Array: TypeAlias = Any
 Device: TypeAlias = Any
 DType: TypeAlias = Any
+
+
+__all__ = [
+    "Array",
+    "SupportsArrayNamespace",
+    "DType",
+    "Device",
+    "HasShape",
+    "Namespace",
+    "NestedSequence",
+    "SupportsBufferProtocol",
+]
+
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/numpy/_aliases.py

@@ -173,5 +173,8 @@ def count_nonzero(
     "pow",
 ]
 __all__ += _aliases.__all__
-
 _all_ignore = ["np", "get_xp"]
+
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/numpy/_info.py

@@ -357,3 +357,10 @@ def devices(self) -> list[Device]:
 
         """
         return ["cpu"]
+
+
+__all__ = ["__array_namespace_info__"]
+
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/numpy/_typing.py

@@ -1,8 +1,5 @@
 from __future__ import annotations
 
-__all__ = ["Array", "DType", "Device"]
-_all_ignore = ["np"]
-
 from typing import TYPE_CHECKING, Any, Literal, TypeAlias
 
 import numpy as np
@@ -24,3 +21,10 @@
 else:
     DType: TypeAlias = np.dtype
     Array: TypeAlias = np.ndarray
+
+__all__ = ["Array", "DType", "Device"]
+_all_ignore = ["np"]
+
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/numpy/fft.py

@@ -20,9 +20,15 @@
 fftshift = get_xp(np)(_fft.fftshift)
 ifftshift = get_xp(np)(_fft.ifftshift)
 
+
 __all__ = ["rfft2", "irfft2", "fft2", "ifft2"]
 __all__ += _fft.__all__
 
+
+def __dir__() -> list[str]:
+    return __all__
+
+
 del get_xp
 del np
 del fft_all