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

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 __future__ import annotations 

9 

10from importlib.metadata import EntryPoints, entry_points 

11import logging 

12import sys 

13from typing import Any 

14 

15from dqm_ml_core.api.processor import Processor 

16 

17logger = logging.getLogger(__name__) 

18 

19 

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

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

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

23 

24 Args: 

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

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

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

28 

29 Returns: 

30 A dictionary mapping plugin names to their loaded classes. 

31 """ 

32 try: 

33 # python 3.10+ 

34 plugin_entry_points: EntryPoints = entry_points(group=plugin_group) 

35 except TypeError: 

36 # Old version for older python version 

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

38 return {} 

39 

40 registry = {} 

41 for v in plugin_entry_points: 

42 # Filter base class registry (not callable) 

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

44 obj = v.load() 

45 if base_class is None or issubclass(obj, base_class): 

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

47 registry[v.name] = obj 

48 else: 

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

50 

51 # return a dict to class builder registry 

52 return registry 

53 

54 

55class PluginLoadedRegistry: 

56 """ 

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

58 

59 Components include: 

60 - Metrics (Processor) 

61 - DataLoaders 

62 - OutputWriters 

63 """ 

64 

65 _metrics_registry: dict[str, type[Processor]] | None = None 

66 _features_registry: dict[str, type[Processor]] | None = None 

67 _gap_registry: dict[str, type[Processor]] | None = None 

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

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

70 

71 @classmethod 

72 def get_metrics_registry(cls) -> dict[str, type[Processor]]: 

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

74 

75 Returns: 

76 A dictionary mapping metric processor names to their classes. 

77 """ 

78 if not cls._metrics_registry: 

79 from dqm_ml_core.api.metrics_processor import MetricsProcessor 

80 

81 cls._metrics_registry = load_registered_plugins("dqm_ml.metrics", MetricsProcessor) 

82 

83 return cls._metrics_registry 

84 

85 @classmethod 

86 def get_features_registry(cls) -> dict[str, type[Processor]]: 

87 """Return the registry of available feature extraction processors. 

88 

89 Returns: 

90 A dictionary mapping feature processor names to their classes. 

91 """ 

92 if not cls._features_registry: 

93 from dqm_ml_core.api.features_processor import FeaturesProcessor 

94 

95 cls._features_registry = load_registered_plugins("dqm_ml.features", FeaturesProcessor) 

96 

97 return cls._features_registry 

98 

99 @classmethod 

100 def get_gap_registry(cls) -> dict[str, type[Processor]]: 

101 """Return the registry of available gap processors. 

102 

103 Returns: 

104 A dictionary mapping gap processor names to their classes. 

105 """ 

106 if not cls._gap_registry: 

107 from dqm_ml_core.api.gap_processor import GapProcessor 

108 

109 cls._gap_registry = load_registered_plugins("dqm_ml.gap", GapProcessor) 

110 

111 return cls._gap_registry 

112 

113 @classmethod 

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

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

116 

117 Returns: 

118 A dictionary mapping data loader names to their classes. 

119 """ 

120 if not cls._dataloaders_registry: 

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

122 return cls._dataloaders_registry 

123 

124 @classmethod 

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

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

127 

128 Returns: 

129 A dictionary mapping output writer names to their classes. 

130 """ 

131 if not cls._outputwriter_registry: 

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

133 

134 return cls._outputwriter_registry