Coverage for packages/dqm-ml/src/dqm_ml/dependency.py: 100%

1"""Dependency management utilities for DQM-ML v2. 

2 

3This module provides functions for handling optional dependencies, 

4displaying version information, and discovering available commands. 

5""" 

6 

7from collections.abc import Generator 

8from contextlib import contextmanager 

9import logging 

10from typing import Any 

11 

12from dqm_ml_core import PluginLoadedRegistry 

13from dqm_ml_core._version_ import version as core_version 

14 

15logger = logging.getLogger(__name__) 

16 

17 

18@contextmanager 

19def optional_dependencies(error: str = "ignore") -> Generator[None, None, None]: 

20 """Context manager for handling optional dependencies. 

21 

22 Args: 

23 error: How to handle missing optional dependencies: 

24 - "ignore": Silently continue (default) 

25 - "warn": Print a warning message 

26 - "raise": Raise the ImportError 

27 

28 Yields: 

29 None - this is a context manager that doesn't provide any value. 

30 """ 

31 assert error in {"raise", "warn", "ignore"} 

32 try: 

33 yield None 

34 except ImportError as e: 

35 if error == "raise": 

36 raise e 

37 if error == "warn": 

38 msg = f'Missing optional dependency "{e.name}". Use pip to install.' 

39 print(f"Warning: {msg}") 

40 

41 

42def display_version(arg_list: list[str] | None = None) -> None: 

43 """Print the DQM-ML core version to stdout. 

44 

45 Args: 

46 arg_list: Unused, provided for CLI compatibility. 

47 """ 

48 print(f"DQM-ML version : {core_version}") 

49 

50 

51def display_list_of(arg_list: list[str] | None = None) -> None: 

52 """Print all registered plugins (metrics, dataloaders, output writers). 

53 

54 Args: 

55 arg_list: Unused, provided for CLI compatibility. 

56 """ 

57 # TODO : we display all but we can filter / use extra parameters 

58 print("Available data metrics_registry") 

59 for key, value in PluginLoadedRegistry.get_metrics_registry().items(): 

60 print(f"- {key} - {value}") 

61 

62 print("Available data features_registry") 

63 for key, value in PluginLoadedRegistry.get_features_registry().items(): 

64 print(f"- {key} - {value}") 

65 

66 print("Available data gap_registry") 

67 for key, value in PluginLoadedRegistry.get_gap_registry().items(): 

68 print(f"- {key} - {value}") 

69 

70 print("Available data loaders") 

71 for key, value in PluginLoadedRegistry.get_dataloaders_registry().items(): 

72 print(f"- {key} - {value}") 

73 

74 print("Available outputs writers") 

75 for key, value in PluginLoadedRegistry.get_outputwriter_registry().items(): 

76 print(f"- {key} - {value}") 

77 

78 

79def get_available_command() -> dict[str, Any]: 

80 """Build a dictionary of available CLI commands. 

81 

82 Returns: 

83 Dict mapping command names to their handler functions. 

84 """ 

85 command_list = {"version": display_version, "list": display_list_of} 

86 optional_dep_mode = "warn" 

87 

88 # We import available command for dqml cli 

89 with optional_dependencies(optional_dep_mode): 

90 from dqm_ml_job._version_ import version as pipeline_version 

91 from dqm_ml_job.cli import execute 

92 

93 logger.debug(f"Different dqm-ml-job version {pipeline_version}") 

94 command_list["process"] = execute 

95 

96 return command_list