Tighten memory detection and target validation

Author: trissim

src/arraybridge/converters.py

@@ -6,7 +6,7 @@
 
 from arraybridge.converters_registry import get_converter
 from arraybridge.framework_config import _FRAMEWORK_CONFIG
-from arraybridge.types import MemoryType
+from arraybridge.types import MemoryType, VALID_MEMORY_TYPES
 
 
 def convert_memory(data: Any, source_type: str, target_type: str, gpu_id: int) -> Any:
@@ -26,6 +26,13 @@ def convert_memory(data: Any, source_type: str, target_type: str, gpu_id: int) -
         ValueError: If source_type or target_type is invalid
         MemoryConversionError: If conversion fails
     """
+    if isinstance(target_type, MemoryType):
+        target_type = target_type.value
+    if target_type not in VALID_MEMORY_TYPES:
+        raise ValueError(
+            f"Invalid target_type '{target_type}'. Available types: {sorted(VALID_MEMORY_TYPES)}"
+        )
+
     converter = get_converter(source_type)  # Will raise ValueError if invalid
     method = getattr(converter, f"to_{target_type}")
     return method(data, gpu_id)
@@ -51,10 +58,14 @@ def detect_memory_type(data: Any) -> str:
     # Check all frameworks using their module names from config
     module_name = type(data).__module__
 
+    top_level = module_name.split(".")[0]
+
     for mem_type, config in _FRAMEWORK_CONFIG.items():
         import_name = config["import_name"]
-        # Check if module name starts with or contains the import name
-        if module_name.startswith(import_name) or import_name in module_name:
+        aliases = {import_name}
+        if import_name == "jax":
+            aliases.add("jaxlib")
+        if top_level in aliases:
             return mem_type.value
 
     raise ValueError(f"Unknown memory type for {type(data)} (module: {module_name})")

src/arraybridge/framework_config.py

@@ -120,17 +120,6 @@ def _tensorflow_validate_dlpack(obj: Any, mod: Any) -> bool:
             f"Clause 88 violation: Cannot infer DLPack capability."
         )
 
-    # Check GPU
-    """Validate TensorFlow DLPack support."""
-    # Check version
-    major, minor = map(int, mod.__version__.split(".")[:2])
-    if major < 2 or (major == 2 and minor < 12):
-        raise RuntimeError(
-            f"TensorFlow {mod.__version__} does not support stable DLPack. "
-            f"Version 2.12.0+ required. "
-            f"Clause 88 violation: Cannot infer DLPack capability."
-        )
-
     # Check GPU
     device_str = obj.device.lower()
     if "gpu" not in device_str:

tests/test_converters.py

@@ -85,8 +85,7 @@ def test_convert_invalid_source_type(self):
     def test_convert_invalid_target_type(self):
         """Test that invalid target type raises error."""
         arr = np.array([1, 2, 3])
-        # This might raise ValueError or AttributeError depending on implementation
-        with pytest.raises((ValueError, AttributeError)):
+        with pytest.raises(ValueError):
             convert_memory(arr, source_type="numpy", target_type="invalid_type", gpu_id=0)
 
     @pytest.mark.torch