Coverage for agentlib/utils/__init__.py: 93%

1""" 

2Module containing all util 

3functions for the agentlib. 

4 

5Most notably, the custom injection enabling 

6dynamic loading of custom models and modules. 

7""" 

8 

9import importlib.util 

10import os 

11import sys 

12from pathlib import Path 

13 

14from .local_broadcast_broker import LocalBroadcastBroker 

15from .local_broker import LocalBroker 

16from .multi_processing_broker import MultiProcessingBroker 

17 

18 

19def custom_injection(config: dict, module_name: str = None): 

20 """ 

21 Function to dynamically load new python files into 

22 the agentlib. Using this, users may use custom modules 

23 oder custom models together with the existing agentlib objects. 

24 

25 Args: 

26 config (dict): Config dict containing the following items: 

27 file (str): Filepath to a python file (.py) 

28 class_name (str): Name of the class to be imported 

29 module_name (str, optional): Name of the imported module 

30 in the sys.modules list. Carefully check if duplicate 

31 module keys raise unexpected behaviour. If so, 

32 use randomly generated strings or similar in classes 

33 calling this function. Default is None. 

34 In that case, the path is converted to a matching string. 

35 

36 Returns: 

37 class (object): The class object specified by class_name 

38 """ 

39 assert "file" in config, ( 

40 "For custom module injection, the config type dict has to " 

41 "contain a 'file'-key with an existing python file as value" 

42 ) 

43 assert "class_name" in config, ( 

44 "For custom module injection, the config type dict has to " 

45 "contain a 'class_name'-key with a string as value " 

46 "specifying the class to inject" 

47 ) 

48 file = config.get("file") 

49 class_name = config.get("class_name") 

50 if not isinstance(file, (str, Path)): 

51 raise TypeError(f"Given file is not a string but {type(file)}") 

52 # Convert to Path object 

53 file = Path(file) 

54 # Check if file is a valid filepath 

55 if not os.path.isfile(file): 

56 raise FileNotFoundError( 

57 f"Given file '{str(file)}' was not found on your device." 

58 ) 

59 

60 # Build module_name if not given: 

61 if module_name is None: 

62 # Build a unique module_name to be imported based on the path 

63 module_name = ".".join([p.name for p in file.parents][:-1] + [file.stem]) 

64 

65 # Custom file import 

66 try: 

67 # Check if the module_name is already present 

68 if module_name in sys.modules: 

69 custom_module = sys.modules[module_name] 

70 else: 

71 spec = importlib.util.spec_from_file_location(module_name, file) 

72 custom_module = importlib.util.module_from_spec(spec) 

73 sys.modules[module_name] = custom_module 

74 spec.loader.exec_module(custom_module) 

75 except ImportError as err: 

76 raise ImportError( 

77 f"Could not inject given module '{class_name}' at '{file}' due to import " 

78 "error. Carefully check for circular imports and partially " 

79 "imported objects based on the following error message: " 

80 f"{err}" 

81 ) from err 

82 try: 

83 return custom_module.__dict__[class_name] 

84 except KeyError: 

85 raise ImportError( 

86 f"Given module '{custom_module}' does not " 

87 f"contain the specified class {class_name}" 

88 )