Coverage for src/arraybridge/converters.py: 96%

1"""Memory conversion public API for OpenHCS.""" 

2 

3from typing import Any, cast 

4 

5import numpy as np 

6 

7from arraybridge.exceptions import MemoryConversionError 

8from arraybridge.types import VALID_MEMORY_TYPES, MemoryType 

9 

10 

11def convert_memory( 

12 data: Any, 

13 source_type: str | MemoryType, 

14 target_type: str | MemoryType, 

15 gpu_id: int, 

16) -> Any: 

17 """ 

18 Convert data between memory types using the unified converter infrastructure. 

19 

20 Args: 

21 data: The data to convert 

22 source_type: The source memory type (e.g., "numpy", "torch") 

23 target_type: The target memory type (e.g., "cupy", "jax") 

24 gpu_id: The target GPU device ID 

25 

26 Returns: 

27 The converted data in the target memory type 

28 

29 Raises: 

30 ValueError: If source_type or target_type is invalid 

31 MemoryConversionError: If conversion fails 

32 """ 

33 source_name = source_type.value if isinstance(source_type, MemoryType) else source_type 

34 target_name = target_type.value if isinstance(target_type, MemoryType) else target_type 

35 if source_name not in VALID_MEMORY_TYPES: 

36 raise ValueError( 

37 f"Invalid source_type '{source_name}'. Available types: {sorted(VALID_MEMORY_TYPES)}" 

38 ) 

39 if target_name not in VALID_MEMORY_TYPES: 

40 raise ValueError( 

41 f"Invalid target_type '{target_name}'. Available types: {sorted(VALID_MEMORY_TYPES)}" 

42 ) 

43 

44 source = MemoryType(source_name) 

45 target = MemoryType(target_name) 

46 try: 

47 return source.convert_to(data, target, gpu_id) 

48 except MemoryConversionError: 

49 raise 

50 except Exception as error: 

51 raise MemoryConversionError( 

52 source_type=source_name, 

53 target_type=target_name, 

54 method="MemoryType.convert_to", 

55 reason=str(error), 

56 ) from error 

57 

58 

59def detect_memory_type(data: Any) -> str: 

60 """ 

61 Detect the memory type of data using framework config. 

62 

63 Args: 

64 data: The data to detect 

65 

66 Returns: 

67 The detected memory type string (e.g., "numpy", "torch") 

68 

69 Raises: 

70 ValueError: If memory type cannot be detected 

71 """ 

72 # NumPy special case (most common, check first) 

73 if isinstance(data, np.ndarray): 

74 return cast(str, MemoryType.NUMPY.value) 

75 

76 # Check all frameworks using their module names from config 

77 module_name = type(data).__module__ 

78 

79 top_level = module_name.split(".")[0] 

80 

81 for mem_type in MemoryType: 

82 if top_level in mem_type.recognized_module_names: 

83 return cast(str, mem_type.value) 

84 

85 raise ValueError(f"Unknown memory type for {type(data)} (module: {module_name})")