Coverage for packages / dqm-ml-core / src / dqm_ml_core / utils / registry.py: 88%

1"""Plugin registry for dynamically loading DQM components. 

2 

3This module contains the PluginLoadedRegistry class and load_registered_plugins 

4function for discovering and loading metric processors, data loaders, 

5and output writers via Python entry points. 

6""" 

7 

8from importlib.metadata import EntryPoints, entry_points 

9import logging 

10import sys 

11from typing import Any 

12 

13from dqm_ml_core import DatametricProcessor 

14 

15logger = logging.getLogger(__name__) 

16 

17 

18# TODO once a base class for all registry created, dict shall have dict[str, base_class] 

19def load_registered_plugins(plugin_group: str, base_class: Any, base_name: str = "default") -> dict[str, Any]: 

20 """Discover and load plugins registered via Python entry points. 

21 

22 Args: 

23 plugin_group: The entry point group name (e.g., 'dqm_ml.metrics'). 

24 base_class: Optional base class to verify plugin type safety. 

25 base_name: Name of the base class to ignore during discovery. 

26 

27 Returns: 

28 A dictionary mapping plugin names to their loaded classes. 

29 """ 

30 try: 

31 # python 3.10+ 

32 plugin_entry_points: EntryPoints = entry_points(group=plugin_group) 

33 except TypeError: 

34 # Old version for older python version 

35 logger.warning(f"Old python version not supported: {sys.version_info}") 

36 return {} 

37 

38 registry = {} 

39 for v in plugin_entry_points: 

40 # Filter base class registry (not callable) 

41 if v.name != base_name: 41 ↛ 39line 41 didn't jump to line 39 because the condition on line 41 was always true

42 obj = v.load() 

43 if base_class is None or issubclass(obj, base_class): 43 ↛ 47line 43 didn't jump to line 47 because the condition on line 43 was always true

44 logger.debug(f"Referencing {plugin_group} - {v.name} class {obj} from {base_class}") 

45 registry[v.name] = obj 

46 else: 

47 logger.error(f"Entry point {plugin_group} - {v.name} class {obj} not derived from {base_class} ignored") 

48 

49 # return a dict to class builder registry 

50 return registry 

51 

52 

53class PluginLoadedRegistry: 

54 """ 

55 Singleton registry that provides lazy access to all registered DQM components. 

56 

57 Components include: 

58 - Metrics (DatametricProcessor) 

59 - DataLoaders 

60 - OutputWriters 

61 """ 

62 

63 _metrics_registry: dict[str, type[DatametricProcessor]] | None = None 

64 _dataloaders_registry: dict[str, Any] | None = None 

65 _outputwriter_registry: dict[str, Any] | None = None 

66 

67 @classmethod 

68 def get_metrics_registry(cls) -> dict[str, type[DatametricProcessor]]: 

69 """Return the registry of available metric processors. 

70 

71 Returns: 

72 A dictionary mapping metric processor names to their classes. 

73 """ 

74 if not cls._metrics_registry: 

75 cls._metrics_registry = load_registered_plugins("dqm_ml.metrics", DatametricProcessor) 

76 

77 return cls._metrics_registry 

78 

79 @classmethod 

80 def get_dataloaders_registry(cls) -> dict[str, Any]: 

81 """Return the registry of available data loaders. 

82 

83 Returns: 

84 A dictionary mapping data loader names to their classes. 

85 """ 

86 if not cls._dataloaders_registry: 

87 cls._dataloaders_registry = load_registered_plugins("dqm_ml.dataloaders", None) # TODO add base class 

88 return cls._dataloaders_registry 

89 

90 @classmethod 

91 def get_outputwriter_registry(cls) -> dict[str, Any]: 

92 """Return the registry of available output writers. 

93 

94 Returns: 

95 A dictionary mapping output writer names to their classes. 

96 """ 

97 if not cls._outputwriter_registry: 

98 cls._outputwriter_registry = load_registered_plugins("dqm_ml.outputwriter", None) # TODO add base class 

99 

100 return cls._outputwriter_registry