Coverage for core / src / sensorkit / common / importutil.py: 98%

1# SPDX-License-Identifier: Apache-2.0 

2"""Utilities for dynamic module loading and object resolution from string specifiers.""" 

3 

4import importlib 

5import importlib.util 

6import inspect 

7import pathlib 

8import sys 

9 

10 

11def get_caller_module(*, depth: int): 

12 """Retrieves the module of the calling function at the specified stack depth.""" 

13 # Account for the direct caller frame. 

14 depth += 1 

15 

16 # Find the requested module by inspecting the call stack. 

17 stack = inspect.stack() 

18 

19 if len(stack) <= depth: 

20 return None 

21 

22 frame = stack[depth][0] 

23 module = inspect.getmodule(frame) 

24 

25 return module 

26 

27 

28def module_from_file(path: pathlib.Path, name: str): 

29 """Load and return a Python module from a filesystem path, handling relative imports within it.""" 

30 # Dynamic load boilerplate. 

31 spec = importlib.util.spec_from_file_location(name, path) 

32 module = importlib.util.module_from_spec(spec) 

33 

34 # We must temporarily add the file's containing directory to the path for any relative imports 

35 # it might contain to work properly. 

36 sys.path.append(str(path.parent)) 

37 

38 try: 

39 # Evaluate the module code. 

40 spec.loader.exec_module(module) 

41 except Exception as e: 

42 # Wrap the exception so a downstream ModuleNotFoundError won't confuse things. 

43 raise ImportError from e 

44 finally: 

45 # Restore sys.path even if import fails. 

46 sys.path.pop() 

47 

48 return module 

49 

50 

51def import_module_or_file(name: str): 

52 """Import a module by name or path, handling relative imports within it.""" 

53 path = pathlib.Path(name) 

54 

55 if path.exists(): 

56 module = module_from_file(path, path.stem) 

57 sys.modules[path.stem] = module 

58 return module 

59 

60 return importlib.import_module(name) 

61 

62 

63def obj_from_spec[T]( 

64 *, 

65 spec: str, 

66 base: type[T], 

67 subclass: bool = False, 

68 load_file: bool = False, 

69) -> T: 

70 """Find and return an object of a given type based on a string specifier.""" 

71 module_name, sep, obj_name = spec.rpartition(":") 

72 

73 if not sep: 

74 module_name = spec 

75 obj_name = None 

76 

77 module = ( 

78 import_module_or_file(module_name) 

79 if load_file 

80 else importlib.import_module(module_name) 

81 ) 

82 haystack = [obj_name] if obj_name else dir(module) 

83 

84 for symbol in haystack: 

85 obj = getattr(module, symbol) 

86 

87 if subclass: 

88 if isinstance(obj, type) and obj is not base and issubclass(obj, base): 

89 return obj 

90 elif isinstance(obj, base): 

91 return obj 

92 

93 kind = "subclass" if subclass else "instance" 

94 raise ValueError(f"no {base.__name__} {kind} matching spec found in {module_name}")