arraybridge 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

arraybridge/__init__.py

@@ -3,37 +3,49 @@ arraybridge: Unified API for NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclespe
 
 This package provides automatic memory type conversion, declarative decorators,
 and unified utilities for working with multiple array/tensor frameworks.
-"""
 
-__version__ = "0.3.2"
+Package-root imports stay lightweight: submodule exports load on first
+attribute access instead of at import time. Declaration-only consumers
+(memory-type tables, configuration modules) therefore do not pay the
+NumPy/numcodecs import cost.
+"""
 
-from . import decorators as _decorators
-from .array_geometry import ArrayGeometry
-from .array_payload import ArrayPayload
-from .converters import convert_memory, detect_memory_type
-from .dtype_scaling import SCALING_FUNCTIONS
-from .exceptions import MemoryConversionError
-from .framework_config import _FRAMEWORK_CONFIG
-from .framework_ops import _FRAMEWORK_OPS
-from .gpu_cleanup import cleanup_all_gpu_frameworks
-from .oom_recovery import _execute_with_oom_recovery
-from .slice_processing import process_slices
-from .stack_utils import stack_slices, unstack_slices
-from .types import (
-    CPU_MEMORY_TYPES,
-    GPU_MEMORY_TYPES,
-    SUPPORTED_MEMORY_TYPES,
-    MemoryContractAttribute,
-    MemoryType,
-)
-from .utils import _ensure_module, _get_device_id, _supports_dlpack
+__version__ = "0.3.4"
 
-DtypeConversion = _decorators.DtypeConversion
-SliceBySliceRuntimeParameter = _decorators.SliceBySliceRuntimeParameter
-memory_types = _decorators.memory_types
-wrap_dtype_preserving_callable = _decorators.wrap_dtype_preserving_callable
-for _memory_type in MemoryType:
-    globals()[_memory_type.value] = getattr(_decorators, _memory_type.value)
+_LAZY_EXPORTS: dict[str, str] = {
+    "MemoryType": ".types",
+    "MemoryContractAttribute": ".types",
+    "ArrayPayload": ".array_payload",
+    "ArrayGeometry": ".array_geometry",
+    "CPU_MEMORY_TYPES": ".types",
+    "GPU_MEMORY_TYPES": ".types",
+    "SUPPORTED_MEMORY_TYPES": ".types",
+    "convert_memory": ".converters",
+    "detect_memory_type": ".converters",
+    "memory_types": ".decorators",
+    "DtypeConversion": ".decorators",
+    "SliceBySliceRuntimeParameter": ".decorators",
+    "wrap_dtype_preserving_callable": ".decorators",
+    "stack_slices": ".stack_utils",
+    "unstack_slices": ".stack_utils",
+    "process_slices": ".slice_processing",
+    "cleanup_all_gpu_frameworks": ".gpu_cleanup",
+    "MemoryConversionError": ".exceptions",
+    "SCALING_FUNCTIONS": ".dtype_scaling",
+    "_FRAMEWORK_CONFIG": ".framework_config",
+    "_FRAMEWORK_OPS": ".framework_ops",
+    "_execute_with_oom_recovery": ".oom_recovery",
+    "_ensure_module": ".utils",
+    "_supports_dlpack": ".utils",
+    "_get_device_id": ".utils",
+    # Decorator exports named after each memory type (numpy, cupy, torch, ...).
+    "numpy": ".decorators",
+    "cupy": ".decorators",
+    "torch": ".decorators",
+    "tensorflow": ".decorators",
+    "jax": ".decorators",
+    "pyclesperanto": ".decorators",
+}
 
 __all__ = [
     # Types
@@ -72,4 +84,26 @@ __all__ = [
     "_ensure_module",
     "_supports_dlpack",
     "_get_device_id",
-] + [memory_type.value for memory_type in MemoryType]
+    "numpy",
+    "cupy",
+    "torch",
+    "tensorflow",
+    "jax",
+    "pyclesperanto",
+]
+
+
+def __getattr__(name: str):
+    module_name = _LAZY_EXPORTS.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    import importlib
+
+    value = getattr(importlib.import_module(module_name, __name__), name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(_LAZY_EXPORTS))

arraybridge/types.py

@@ -5,6 +5,8 @@ This module defines the MemoryType enum and related constants for managing
 different array/tensor frameworks.
 """
 
+from __future__ import annotations
+
 import importlib
 import importlib.util
 import logging
@@ -14,17 +16,11 @@ from collections.abc import Callable, Iterator, Mapping, MutableMapping
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, TypeVar, cast
-
-from arraybridge.array_operations import (
-    CUPY_OPERATIONS,
-    JAX_OPERATIONS,
-    NUMPY_OPERATIONS,
-    PYCLESPERANTO_OPERATIONS,
-    TENSORFLOW_OPERATIONS,
-    TORCH_OPERATIONS,
-    ArrayOperations,
-)
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, TypeVar, cast
+
+if TYPE_CHECKING:
+    from arraybridge.array_operations import ArrayOperations
 
 T = TypeVar("T")
 logger = logging.getLogger(__name__)
@@ -41,6 +37,55 @@ StreamScopeFactory = Callable[[Any, Any], AbstractContextManager[None]]
 DLPackExporter = Callable[[Any, Any], Any | None]
 DLPackValidator = Callable[[Any, Any], bool]
 OOMMatcher = Callable[[BaseException, Any | None], bool]
+SubprocessEnvironmentResolver = Callable[[Mapping[str, str]], dict[str, str]]
+
+
+def _identity_subprocess_environment(
+    environment: Mapping[str, str],
+) -> dict[str, str]:
+    return dict(environment)
+
+
+def _nvidia_wheel_library_paths() -> tuple[str, ...]:
+    """Return native-library directories declared by installed NVIDIA wheels."""
+
+    try:
+        spec = importlib.util.find_spec("nvidia")
+    except (ImportError, ModuleNotFoundError, ValueError):
+        return ()
+    if spec is None or spec.submodule_search_locations is None:
+        return ()
+
+    paths: set[str] = set()
+    for package_root in map(Path, spec.submodule_search_locations):
+        try:
+            components = tuple(package_root.iterdir())
+        except OSError:
+            continue
+        for component in components:
+            for library_directory_name in ("lib", "bin"):
+                candidate = component / library_directory_name
+                if candidate.is_dir():
+                    paths.add(str(candidate))
+    return tuple(sorted(paths))
+
+
+def _nvidia_wheel_subprocess_environment(
+    environment: Mapping[str, str],
+) -> dict[str, str]:
+    """Prepend installed NVIDIA wheel libraries to a child environment."""
+
+    prepared = dict(environment)
+    library_paths = _nvidia_wheel_library_paths()
+    if not library_paths:
+        return prepared
+
+    search_variable = "PATH" if os.name == "nt" else "LD_LIBRARY_PATH"
+    existing_paths = tuple(
+        path for path in prepared.get(search_variable, "").split(os.pathsep) if path
+    )
+    prepared[search_variable] = os.pathsep.join(dict.fromkeys((*library_paths, *existing_paths)))
+    return prepared
 
 
 class MemoryContractAttribute(str, Enum):
@@ -501,6 +546,7 @@ class FrameworkRuntime:
     dlpack_exporter: DLPackExporter | None = None
     dlpack_validator: DLPackValidator = _protocol_dlpack
     oom_matcher: OOMMatcher = _never_oom
+    subprocess_environment: SubprocessEnvironmentResolver = _identity_subprocess_environment
 
 
 class _MemoryTypeFields:
@@ -510,7 +556,7 @@ class _MemoryTypeFields:
     module_aliases: tuple[str, ...]
     import_environment: tuple[tuple[str, str], ...]
     _runtime: FrameworkRuntime
-    _operations: ArrayOperations
+    _operations_name: str
 
 
 class MemoryType(_MemoryTypeFields, Enum):
@@ -520,7 +566,7 @@ class MemoryType(_MemoryTypeFields, Enum):
         cls,
         value: str,
         *declaration: Any,
-    ) -> "MemoryType":
+    ) -> MemoryType:
         (
             import_name,
             display_name,
@@ -538,9 +584,18 @@ class MemoryType(_MemoryTypeFields, Enum):
         member.module_aliases = cast(tuple[str, ...], module_aliases)
         member.import_environment = cast(tuple[tuple[str, str], ...], import_environment)
         member._runtime = cast(FrameworkRuntime, runtime)
-        member._operations = cast(ArrayOperations, operations)
+        member._operations_name = cast(str, operations)
         return member
 
+    @property
+    def _operations(self) -> ArrayOperations:
+        """Resolve this declaration's operations registry on first access."""
+
+        from arraybridge import array_operations
+
+        operations: ArrayOperations = getattr(array_operations, self._operations_name)
+        return operations
+
     NUMPY = (
         "numpy",
         "numpy",
@@ -549,7 +604,7 @@ class MemoryType(_MemoryTypeFields, Enum):
         (),
         (),
         FrameworkRuntime(oom_matcher=_numpy_oom),
-        NUMPY_OPERATIONS,
+        "NUMPY_OPERATIONS",
     )
     CUPY = (
         "cupy",
@@ -571,8 +626,9 @@ class MemoryType(_MemoryTypeFields, Enum):
             dlpack_importer=_cupy_from_dlpack,
             dlpack_exporter=_protocol_dlpack_export,
             oom_matcher=_cupy_oom,
+            subprocess_environment=_nvidia_wheel_subprocess_environment,
         ),
-        CUPY_OPERATIONS,
+        "CUPY_OPERATIONS",
     )
     TORCH = (
         "torch",
@@ -595,7 +651,7 @@ class MemoryType(_MemoryTypeFields, Enum):
             dlpack_exporter=_protocol_dlpack_export,
             oom_matcher=_torch_oom,
         ),
-        TORCH_OPERATIONS,
+        "TORCH_OPERATIONS",
     )
     TENSORFLOW = (
         "tensorflow",
@@ -614,7 +670,7 @@ class MemoryType(_MemoryTypeFields, Enum):
             dlpack_validator=_tensorflow_dlpack,
             oom_matcher=_tensorflow_oom,
         ),
-        TENSORFLOW_OPERATIONS,
+        "TENSORFLOW_OPERATIONS",
     )
     JAX = (
         "jax",
@@ -632,7 +688,7 @@ class MemoryType(_MemoryTypeFields, Enum):
             dlpack_exporter=_protocol_dlpack_export,
             oom_matcher=_jax_oom,
         ),
-        JAX_OPERATIONS,
+        "JAX_OPERATIONS",
     )
     PYCLESPERANTO = (
         "pyclesperanto",
@@ -649,7 +705,7 @@ class MemoryType(_MemoryTypeFields, Enum):
             move_to_active_device=_pyclesperanto_device_move,
             oom_matcher=_pyclesperanto_oom,
         ),
-        PYCLESPERANTO_OPERATIONS,
+        "PYCLESPERANTO_OPERATIONS",
     )
 
     @property
@@ -682,6 +738,20 @@ class MemoryType(_MemoryTypeFields, Enum):
         for name, value in self.import_environment:
             os.environ.setdefault(name, value)
 
+    @classmethod
+    def subprocess_environment(
+        cls,
+        environment: Mapping[str, str] | None = None,
+    ) -> dict[str, str]:
+        """Project framework import requirements into a child environment."""
+
+        prepared = dict(os.environ if environment is None else environment)
+        for memory_type in cls:
+            for name, value in memory_type.import_environment:
+                prepared.setdefault(name, value)
+            prepared = memory_type._runtime.subprocess_environment(prepared)
+        return prepared
+
     def loaded_module(self) -> Any | None:
         """Return an already-loaded framework without causing an import."""
 
@@ -887,7 +957,7 @@ class MemoryType(_MemoryTypeFields, Enum):
         payload = data if isinstance(data, DLPackPayload) else DLPackPayload(data, data)
         return importer(payload, framework)
 
-    def convert_to(self, data: Any, target: "MemoryType", device_id: int) -> Any:
+    def convert_to(self, data: Any, target: MemoryType, device_id: int) -> Any:
         """Convert one array through the source and target declarations."""
 
         if self is target:

{arraybridge-0.3.2.dist-info → arraybridge-0.3.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arraybridge
-Version: 0.3.2
+Version: 0.3.4
 Summary: Unified API for NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclesperanto with automatic memory type conversion
 Project-URL: Homepage, https://github.com/OpenHCSDev/arraybridge
 Project-URL: Documentation, https://arraybridge.readthedocs.io

{arraybridge-0.3.2.dist-info → arraybridge-0.3.4.dist-info}/RECORD

@@ -1,4 +1,4 @@
-arraybridge/__init__.py,sha256=xFy8ESRzsIy8ydjdPQs33C5Ed6HH6YUsn65HSmQTtzM,2326
+arraybridge/__init__.py,sha256=a27ES1ItSBsT77vP3f12ILNu0MzrvQfcpCBVDXwi-Sg,3184
 arraybridge/array_geometry.py,sha256=CopfGmFVKwV5l-lOTImhm1YVh18weUng0ReEF58qlPI,1689
 arraybridge/array_operations.py,sha256=fp69o9BYcCzfhmLcWrH68SVGSfPL3nb6Cty-Amd_oQM,13082
 arraybridge/array_payload.py,sha256=6_jJwPRya0MHrUNSA9kXPlTISpLa0G74rqDValGXKx4,1079
@@ -13,9 +13,9 @@ arraybridge/gpu_cleanup.py,sha256=ADI4wMLk9GIO9rJ6BJuNk3uZIyANoAqJ7VFTnLJzrbE,25
 arraybridge/oom_recovery.py,sha256=hGvOuJHVSopiIqblzoMFlqW8_oK6ZlAGF_l4pfEljeI,2403
 arraybridge/slice_processing.py,sha256=5R10OmHW8BSbJFwT1URd2Sz3P15o5EhaI2VnVUUIvkQ,4450
 arraybridge/stack_utils.py,sha256=FVikN0XG-Xryw-1p9MaSRUTZet7kLcy0gTEERSjc9r4,6915
-arraybridge/types.py,sha256=1otpHqFJfGc9RY40VCzNCkqpTnR8cfckhY44HnuGudI,33191
+arraybridge/types.py,sha256=e5R2DVflQAjXrfTk033KRdiiwoip3IQAvgDCYUjPdiQ,35807
 arraybridge/utils.py,sha256=RqThtYScPEOEQG491dA8b_VvUTjgcFhwLbh0erQu5Ns,4833
-arraybridge-0.3.2.dist-info/METADATA,sha256=tX2dNiYptc7DpXQUqV8x5BX5RzVSMLOD9wM-4xQmg5o,5275
-arraybridge-0.3.2.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
-arraybridge-0.3.2.dist-info/licenses/LICENSE,sha256=xagEoeTAj1WT64RmyR3E6HH-eTGdgXN6gqPMUUt7L_Y,1070
-arraybridge-0.3.2.dist-info/RECORD,,
+arraybridge-0.3.4.dist-info/METADATA,sha256=jq9-ULcNwbZo7cMtZlhhqj4F2QlV-8PrpdM1Tv0tLCE,5275
+arraybridge-0.3.4.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+arraybridge-0.3.4.dist-info/licenses/LICENSE,sha256=xagEoeTAj1WT64RmyR3E6HH-eTGdgXN6gqPMUUt7L_Y,1070
+arraybridge-0.3.4.dist-info/RECORD,,