Coverage for openhcs/utils/import_utils.py: 81.0%

1""" 

2Utility functions for handling optional imports in OpenHCS. 

3 

4This module provides functions for importing optional dependencies 

5in a way that allows functions to be registered even if their 

6dependencies aren't available at registration time. 

7""" 

8import importlib 

9from typing import Optional, Any, Type 

10 

11 

12def optional_import(module_name: str) -> Optional[Any]: 

13 """ 

14 Import a module if available, otherwise return None. 

15 

16 Args: 

17 module_name: Name of the module to import (can be a dotted path) 

18 

19 Returns: 

20 The imported module if available, None otherwise 

21 """ 

22 try: 

23 return importlib.import_module(module_name) 

24 except (ImportError, ModuleNotFoundError, AttributeError): # Added AttributeError for safety 

25 return None 

26 

27 

28def create_placeholder_class(name: str, base_class: Optional[Any] = None, 

29 required_library: str = "") -> Type: 

30 """ 

31 Create a placeholder class when a required library is not available. 

32 

33 This function generates a placeholder class that can be used in place of a class 

34 that depends on an optional library. The placeholder class will raise an ImportError 

35 when any of its methods are called or attributes are accessed (excluding __init__, __name__, __doc__). 

36 

37 Args: 

38 name: Name of the class to be created. 

39 base_class: Optional base class to inherit from if the actual library is available. 

40 If the library and thus the base_class are None, a placeholder is created. 

41 required_library: Name of the required library (for error messages). 

42 

43 Returns: 

44 Either the base_class itself if it's not None (meaning the library was available), 

45 or a newly created placeholder class that raises ImportError on use. 

46 

47 Example: 

48 ```python 

49 torch = optional_import("torch") 

50 nn = optional_import("torch.nn") if torch else None 

51 

52 # If nn and nn.Module are available, Module will be nn.Module 

53 # Otherwise, Module will be a placeholder. 

54 Module = create_placeholder_class( 

55 "Module", 

56 base_class=nn.Module if nn else None, 

57 required_library="PyTorch" 

58 ) 

59 

60 class MyModel(Module): # Inherits from nn.Module or Placeholder 

61 def __init__(self): 

62 super().__init__() # Works for both 

63 if isinstance(self, nn.Module): # Check if real or placeholder 

64 self.linear = nn.Linear(10,1) # Only if real 

65  

66 def forward(self, x): 

67 # This would raise ImportError if Module is a placeholder and self.linear wasn't set 

68 # or if super().forward() was called on a placeholder. 

69 if isinstance(self, nn.Module): 

70 return self.linear(x) 

71 else: 

72 # Placeholder specific behavior or raise error 

73 raise ImportError(f"PyTorch is required to use MyModel.forward") 

74 

75 ``` 

76 """ 

77 if base_class is not None: 77 ↛ 80line 77 didn't jump to line 80 because the condition on line 77 was never true

78 # The library and base class are available, return the actual base class. 

79 # The calling code will then define a class that inherits from this real base. 

80 return base_class 

81 else: 

82 # Create a placeholder class 

83 # This class will be used as a base for classes in other modules. 

84 # When methods of those derived classes are called (especially super().__init__ or super().method()), 

85 # or attributes are accessed, __getattr__ on this placeholder will be triggered if not found. 

86 class Placeholder: 

87 _required_library_name = required_library or "An optional library" 

88 

89 def __init__(self, *args: Any, **kwargs: Any) -> None: 

90 # The __init__ of a placeholder should generally do nothing or 

91 # just store args/kwargs if needed for some very specific placeholder logic. 

92 # It should NOT try to call super().__init__ if it's meant to be a root placeholder. 

93 # If this placeholder *itself* is meant to inherit from something, that's a different pattern. 

94 # For replacing e.g. nn.Module, this __init__ is fine. 

95 pass 

96 

97 def __getattr__(self, item: str) -> Any: 

98 # This will be called for any attribute not found on the instance. 

99 # This includes methods. 

100 raise ImportError( 

101 f"{self._required_library_name} is required to use the attribute/method '{item}' " 

102 f"of class '{name}' (which is a placeholder)." 

103 ) 

104 

105 Placeholder.__name__ = name 

106 Placeholder.__doc__ = ( 

107 f"Placeholder for '{name}' when {required_library or 'its required library'} " 

108 "is not available. Accessing attributes or methods will raise an ImportError." 

109 ) 

110 return Placeholder