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

arraybridge/__init__.py

@@ -5,7 +5,7 @@ This package provides automatic memory type conversion, declarative decorators,
 and unified utilities for working with multiple array/tensor frameworks.
 """
 
-__version__ = "0.3.2"
+__version__ = "0.3.3"
 
 from . import decorators as _decorators
 from .array_geometry import ArrayGeometry

arraybridge/types.py

@@ -14,6 +14,7 @@ from collections.abc import Callable, Iterator, Mapping, MutableMapping
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
 from enum import Enum
+from pathlib import Path
 from typing import Any, TypeVar, cast
 
 from arraybridge.array_operations import (
@@ -41,6 +42,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 +551,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:
@@ -571,6 +622,7 @@ 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,
     )
@@ -682,6 +734,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."""
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arraybridge
-Version: 0.3.2
+Version: 0.3.3
 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.3.dist-info}/RECORD

@@ -1,4 +1,4 @@
-arraybridge/__init__.py,sha256=xFy8ESRzsIy8ydjdPQs33C5Ed6HH6YUsn65HSmQTtzM,2326
+arraybridge/__init__.py,sha256=TmDOlpJsQxwY_k9aQ6YlGav2mZLLPzmh2tWbAOY1FyQ,2326
 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=F_aJLJnZ46KBmCziHjtfkkvpqzHkc-kDjbsuRc98uak,35588
 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.3.dist-info/METADATA,sha256=eloa3YEUoJnRo0S3n6JgClhYWy8J8sdAtINIwsroOf8,5275
+arraybridge-0.3.3.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+arraybridge-0.3.3.dist-info/licenses/LICENSE,sha256=xagEoeTAj1WT64RmyR3E6HH-eTGdgXN6gqPMUUt7L_Y,1070
+arraybridge-0.3.3.dist-info/RECORD,,