Reorganize __all__

Author: mdboom

cuda_core/cuda/core/__init__.py

@@ -69,90 +69,51 @@ class _PatchedProperty(metaclass=_PatchedPropMeta):
 
 
 from cuda.core import checkpoint, system, utils
-from cuda.core._context import Context, ContextOptions
-from cuda.core._device import Device
-from cuda.core._device_resources import (
-    DeviceResources,
-    SMResource,
-    SMResourceOptions,
-    WorkqueueResource,
-    WorkqueueResourceOptions,
+from cuda.core._context import *
+from cuda.core._context import __all__ as _context_all
+from cuda.core._device import *
+from cuda.core._device import __all__ as _device_all
+from cuda.core._device_resources import *
+from cuda.core._device_resources import __all__ as _device_resources_all
+from cuda.core._event import *
+from cuda.core._event import __all__ as _event_all
+from cuda.core._graphics import *
+from cuda.core._graphics import __all__ as _graphics_all
+from cuda.core._host import *
+from cuda.core._host import __all__ as _host_all
+from cuda.core._launch_config import *
+from cuda.core._launch_config import __all__ as _launch_config_all
+from cuda.core._launcher import *
+from cuda.core._launcher import __all__ as _launcher_all
+from cuda.core._linker import *
+from cuda.core._linker import __all__ as _linker_all
+from cuda.core._memory import *
+from cuda.core._memory import __all__ as _memory_all
+from cuda.core._module import *
+from cuda.core._module import __all__ as _module_all
+from cuda.core._program import *
+from cuda.core._program import __all__ as _program_all
+from cuda.core._stream import *
+from cuda.core._stream import __all__ as _stream_all
+from cuda.core._tensor_map import *
+from cuda.core._tensor_map import __all__ as _tensor_map_all
+
+__all__ = (
+    _context_all
+    + _device_all
+    + _device_resources_all
+    + _event_all
+    + _graphics_all
+    + _host_all
+    + _launch_config_all
+    + _launcher_all
+    + _linker_all
+    + _memory_all
+    + _module_all
+    + _program_all
+    + _stream_all
+    + _tensor_map_all
 )
-from cuda.core._event import Event, EventOptions
-from cuda.core._graphics import GraphicsResource
-from cuda.core._host import Host
-from cuda.core._launch_config import LaunchConfig
-from cuda.core._launcher import launch
-from cuda.core._linker import Linker, LinkerOptions
-from cuda.core._memory import (
-    Buffer,
-    DeviceMemoryResource,
-    DeviceMemoryResourceOptions,
-    GraphMemoryResource,
-    LegacyPinnedMemoryResource,
-    ManagedBuffer,
-    ManagedMemoryResource,
-    ManagedMemoryResourceOptions,
-    MemoryResource,
-    PinnedMemoryResource,
-    PinnedMemoryResourceOptions,
-    VirtualMemoryResource,
-    VirtualMemoryResourceOptions,
-)
-from cuda.core._module import Kernel, ObjectCode
-from cuda.core._program import Program, ProgramOptions
-from cuda.core._stream import (
-    LEGACY_DEFAULT_STREAM,
-    PER_THREAD_DEFAULT_STREAM,
-    Stream,
-    StreamOptions,
-)
-from cuda.core._tensor_map import TensorMapDescriptor, TensorMapDescriptorOptions
-
-# Flat public API of the ``cuda.core`` namespace. Submodule namespaces
-# (``checkpoint``, ``graph``, ``system``, ``texture``, ``utils``) carry their
-# own ``__all__`` and are intentionally not re-listed here.
-__all__ = [
-    "LEGACY_DEFAULT_STREAM",
-    "PER_THREAD_DEFAULT_STREAM",
-    "Buffer",
-    "Context",
-    "ContextOptions",
-    "Device",
-    "DeviceMemoryResource",
-    "DeviceMemoryResourceOptions",
-    "DeviceResources",
-    "Event",
-    "EventOptions",
-    "GraphMemoryResource",
-    "GraphicsResource",
-    "Host",
-    "Kernel",
-    "LaunchConfig",
-    "LegacyPinnedMemoryResource",
-    "Linker",
-    "LinkerOptions",
-    "ManagedBuffer",
-    "ManagedMemoryResource",
-    "ManagedMemoryResourceOptions",
-    "MemoryResource",
-    "ObjectCode",
-    "PinnedMemoryResource",
-    "PinnedMemoryResourceOptions",
-    "Program",
-    "ProgramOptions",
-    "SMResource",
-    "SMResourceOptions",
-    "Stream",
-    "StreamOptions",
-    "TensorMapDescriptor",
-    "TensorMapDescriptorOptions",
-    "VirtualMemoryResource",
-    "VirtualMemoryResourceOptions",
-    "WorkqueueResource",
-    "WorkqueueResourceOptions",
-    "launch",
-]
 
 # isort: split
 # Texture/surface types live under the cuda.core.texture namespace (not the

cuda_core/cuda/core/_device.pyi

@@ -1023,4 +1023,5 @@ class Device:
         .. versionadded:: 1.1.0
         """
 _tls = threading.local()
-_lock = threading.Lock()
+_lock = threading.Lock()
+__all__ = ['Device']

cuda_core/cuda/core/_device.pyx

@@ -61,6 +61,8 @@ _tls = threading.local()
 _lock = threading.Lock()
 cdef bint _is_cuInit = False
 
+__all__ = ['Device']
+
 
 cdef class DeviceProperties:
     """

cuda_core/cuda/core/_event.pyi

@@ -180,6 +180,7 @@ class IPCEventDescriptor:
 
     def __reduce__(self) -> tuple[object, ...]:
         ...
+__all__ = ['Event', 'EventOptions']
 
 def _reduce_event(event: Event) -> tuple[object, ...]:
     ...

cuda_core/cuda/core/_event.pyx

@@ -43,6 +43,8 @@ if TYPE_CHECKING:
     import cuda.bindings.driver  # no-cython-lint
     from cuda.core._device import Device
 
+__all__ = ['Event', 'EventOptions']
+
 
 @dataclass
 cdef class EventOptions:

cuda_core/cuda/core/_host.py

@@ -6,6 +6,8 @@
 import threading
 from typing import ClassVar
 
+__all__ = ["Host"]
+
 
 class Host:
     """Host (CPU) location for managed-memory operations.

cuda_core/cuda/core/_launch_config.pyi

@@ -66,6 +66,7 @@ class LaunchConfig:
     def __hash__(self) -> int:
         ...
 _LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative')
+__all__ = ['LaunchConfig']
 
 def _to_native_launch_config(config: LaunchConfig) -> object:
     """Convert LaunchConfig to native driver CUlaunchConfig.

cuda_core/cuda/core/_launch_config.pyx

@@ -15,6 +15,8 @@ from cuda.core._utils.cuda_utils import (
 
 _LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative')
 
+__all__ = ['LaunchConfig']
+
 
 cdef class LaunchConfig:
     """Customizable launch options.

cuda_core/cuda/core/_launcher.pyi

@@ -8,6 +8,7 @@ from cuda.core._stream import Stream
 from cuda.core.graph import GraphBuilder
 from cuda.core.typing import IsStreamType
 
+__all__ = ['launch']
 
 def launch(stream: Stream | GraphBuilder | IsStreamType, config: LaunchConfig, kernel: Kernel, *kernel_args) -> None:
     """Launches a :obj:`~_module.Kernel`

cuda_core/cuda/core/_launcher.pyx

@@ -24,6 +24,8 @@ if TYPE_CHECKING:
     from cuda.core.graph import GraphBuilder
     from cuda.core.typing import IsStreamType
 
+__all__ = ['launch']
+
 
 def launch(
     stream: Stream | GraphBuilder | IsStreamType,