Coverage for agentlib/modules/__init__.py: 86%

1""" 

2Package containing all modules used by agents. 

3Use the helper functions get_module_type 

4to load module classes from this package. 

5""" 

6 

7from typing import Union, Dict, Iterable, List 

8 

9from agentlib.utils import plugin_import 

10from agentlib.utils.fuzzy_matching import fuzzy_match 

11 

12# Global modules: 

13_MODULE_TYPES = plugin_import.SaveUpdateDict() 

14_LOADED_CORE_MODULES = False 

15 

16 

17def _import_modules(): 

18 """ 

19 Import the module information, 

20 but only if it has not been loaded yet. 

21 """ 

22 global _MODULE_TYPES 

23 if not _LOADED_CORE_MODULES: 

24 _load_core_modules() 

25 

26 # Only return copy, never allow changes in the private object 

27 return _MODULE_TYPES.copy() 

28 

29 

30def _load_core_modules() -> None: 

31 global _LOADED_CORE_MODULES, _MODULE_TYPES 

32 # Communicator 

33 import agentlib.modules.communicator as communicator 

34 

35 _MODULE_TYPES.update(communicator.MODULE_TYPES) 

36 # Controller 

37 import agentlib.modules.controller as controller 

38 

39 _MODULE_TYPES.update(controller.MODULE_TYPES) 

40 # Utils 

41 import agentlib.modules.utils as utils 

42 

43 _MODULE_TYPES.update(utils.MODULE_TYPES) 

44 # Simulator 

45 import agentlib.modules.simulation as simulation 

46 

47 _MODULE_TYPES.update(simulation.MODULE_TYPES) 

48 _LOADED_CORE_MODULES = True 

49 

50 

51def get_module_type(module_type) -> Union[Dict, Iterable]: 

52 """ 

53 Return and load the given module type 

54 

55 Args: 

56 module_type str: 

57 The string identifier to load the module. 

58 

59 Returns: 

60 module BaseModuleType: 

61 The module specified by the given module_type 

62 """ 

63 # Check if it's a plugin 

64 if "." in module_type: 

65 plugin_import.load_plugin( 

66 name=module_type.split(".")[0], 

67 loaded_classes=_MODULE_TYPES, 

68 plugin_types_name="MODULE_TYPES", 

69 ) 

70 # Load the core and plugin modules 

71 module_types = _import_modules() 

72 if module_type in module_types: 

73 return module_types[module_type].import_class() 

74 # Raise error if still here 

75 

76 matches = fuzzy_match(target=module_type, choices=module_types.keys()) 

77 msg = ( 

78 f"Given module_type '{module_type}' is neither in the AgentLib nor in " 

79 f"installed plugins. " 

80 ) 

81 if matches: 

82 msg += f"Did you mean one of these? {', '.join(matches)}" 

83 

84 # Extract names only 

85 raise ModuleNotFoundError(msg) 

86 

87 

88def get_all_module_types(plugins: List[str] = None): 

89 """ 

90 Returns all available module types 

91 

92 Args: 

93 plugins List[str]: 

94 A list of strings being the 

95 plugins to consider in the search. 

96 

97 Returns: 

98 dict: Module types, with the key as the types name 

99 and the value being the ModuleImport instance 

100 """ 

101 for plugin in plugins: 

102 plugin_import.load_plugin( 

103 name=plugin, loaded_classes=_MODULE_TYPES, plugin_types_name="MODULE_TYPES" 

104 ) 

105 module_types = _import_modules() 

106 return module_types