Coverage for agentlib/models/__init__.py: 40%

1"""Package with available models for the agentlib""" 

2 

3from agentlib.utils import plugin_import 

4from agentlib.utils.fuzzy_matching import fuzzy_match 

5from agentlib.utils.plugin_import import ModuleImport 

6 

7UNINSTALLED_MODEL_TYPES = {} 

8 

9MODEL_TYPES = plugin_import.SaveUpdateDict( 

10 **{ 

11 "statespace": ModuleImport( 

12 import_path="agentlib.models.scipy_model", class_name="ScipyStateSpaceModel" 

13 ), 

14 "fmu": ModuleImport( 

15 import_path="agentlib.models.fmu_model", class_name="FmuModel" 

16 ), 

17 } 

18) 

19 

20 

21def get_model_type(model_type) -> type: 

22 """ 

23 Return and load the given module type 

24 

25 Args: 

26 model_type str: 

27 The string identifier to load the module. 

28 

29 Returns: 

30 module BaseModelType: 

31 """ 

32 # Check if it's a plugin 

33 if "." in model_type: 

34 plugin_import.load_plugin( 

35 name=model_type.split(".")[0], 

36 loaded_classes=MODEL_TYPES, 

37 plugin_types_name="MODEL_TYPES", 

38 ) 

39 # Load the core and plugin modules 

40 if model_type in MODEL_TYPES: 

41 return MODEL_TYPES[model_type].import_class() 

42 

43 matches = fuzzy_match(target=model_type, choices=MODEL_TYPES.keys()) 

44 msg = ( 

45 f"Given model_type '{model_type}' is neither in the AgentLib nor in " 

46 f"installed plugins. " 

47 ) 

48 if matches: 

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

50 raise ModuleNotFoundError(msg)