Coverage for ml_workbench/cli_config.py: 17%

1""" 

2CLI tool for processing YAML configuration files. 

3 

4This tool loads a YAML configuration file, processes any include directives, 

5performs variable interpolation, and outputs the resulting merged YAML. 

6""" 

7 

8from __future__ import annotations 

9 

10import argparse 

11import json 

12import sys 

13from pathlib import Path 

14from typing import Optional 

15 

16import yaml 

17 

18from ml_workbench.config import YamlConfig 

19 

20 

21def build_arg_parser() -> argparse.ArgumentParser: 

22 """ 

23 Build the argument parser for the CLI. 

24  

25 Returns 

26 ------- 

27 argparse.ArgumentParser 

28 Configured argument parser 

29 """ 

30 parser = argparse.ArgumentParser( 

31 description="Process YAML configuration files with includes and variable interpolation", 

32 formatter_class=argparse.RawDescriptionHelpFormatter, 

33 epilog=""" 

34Examples: 

35 # Process a YAML file and output the result 

36 %(prog)s config.yaml 

37  

38 # Process with variable substitution 

39 %(prog)s config.yaml --var catalog=prod --var env=production 

40  

41 # Output as JSON instead of YAML 

42 %(prog)s config.yaml --json 

43  

44 # Output to a file 

45 %(prog)s config.yaml -o output.yaml 

46  

47 # Show only specific sections 

48 %(prog)s config.yaml --section datasets 

49 %(prog)s config.yaml --section features --section datasets 

50 """, 

51 ) 

52 

53 parser.add_argument( 

54 "yaml", 

55 type=str, 

56 help="Path to the YAML configuration file", 

57 ) 

58 

59 parser.add_argument( 

60 "-o", 

61 "--output", 

62 type=str, 

63 help="Output file path (default: stdout)", 

64 ) 

65 

66 parser.add_argument( 

67 "--json", 

68 action="store_true", 

69 help="Output as JSON instead of YAML", 

70 ) 

71 

72 parser.add_argument( 

73 "--var", 

74 action="append", 

75 dest="vars", 

76 metavar="KEY=VALUE", 

77 help="Variable for interpolation (can be specified multiple times)", 

78 ) 

79 

80 parser.add_argument( 

81 "--section", 

82 action="append", 

83 dest="sections", 

84 metavar="SECTION", 

85 help="Output only specific section(s) (can be specified multiple times)", 

86 ) 

87 

88 parser.add_argument( 

89 "--no-impute", 

90 action="store_true", 

91 help="Skip dataset type/format imputation", 

92 ) 

93 

94 parser.add_argument( 

95 "--show-variables", 

96 action="store_true", 

97 help="Show resolved variables", 

98 ) 

99 

100 parser.add_argument( 

101 "--strict", 

102 action="store_true", 

103 default=True, 

104 help="Strict mode: fail on missing variables (default: True)", 

105 ) 

106 

107 parser.add_argument( 

108 "--no-strict", 

109 action="store_false", 

110 dest="strict", 

111 help="Non-strict mode: leave unknown placeholders unchanged", 

112 ) 

113 

114 return parser 

115 

116 

117def parse_kv_pairs(pairs: list[str]) -> dict[str, str]: 

118 """ 

119 Parse key=value pairs from command line arguments. 

120  

121 Parameters 

122 ---------- 

123 pairs : list[str] 

124 List of "key=value" strings 

125  

126 Returns 

127 ------- 

128 dict[str, str] 

129 Dictionary of parsed key-value pairs 

130  

131 Raises 

132 ------ 

133 ValueError 

134 If a pair is not in "key=value" format 

135 """ 

136 result = {} 

137 for pair in pairs: 

138 if "=" not in pair: 

139 raise ValueError(f"Invalid key=value pair: {pair}") 

140 key, value = pair.split("=", 1) 

141 result[key.strip()] = value.strip() 

142 return result 

143 

144 

145def filter_sections(data: dict, sections: list[str]) -> dict: 

146 """ 

147 Filter configuration to include only specified sections. 

148  

149 Parameters 

150 ---------- 

151 data : dict 

152 Full configuration dictionary 

153 sections : list[str] 

154 List of section names to include 

155  

156 Returns 

157 ------- 

158 dict 

159 Filtered configuration with only specified sections 

160 """ 

161 result = {} 

162 for section in sections: 

163 if section in data: 

164 result[section] = data[section] 

165 else: 

166 print(f"Warning: Section '{section}' not found in configuration", file=sys.stderr) 

167 return result 

168 

169 

170def main(argv: Optional[list[str]] = None) -> int: 

171 """ 

172 Main entry point for the CLI. 

173  

174 Parameters 

175 ---------- 

176 argv : Optional[list[str]] 

177 Command line arguments (default: sys.argv) 

178  

179 Returns 

180 ------- 

181 int 

182 Exit code (0 for success, non-zero for error) 

183 """ 

184 parser = build_arg_parser() 

185 args = parser.parse_args(argv) 

186 

187 try: 

188 # Parse variable substitutions 

189 variables = parse_kv_pairs(list(args.vars)) if args.vars else {} 

190 

191 # Load configuration 

192 cfg = YamlConfig(args.yaml, strict=args.strict, **variables) 

193 

194 # Get the processed data 

195 data = cfg.to_dict() 

196 

197 # Filter sections if specified 

198 if args.sections: 

199 data = filter_sections(data, args.sections) 

200 

201 # Show variables if requested 

202 if args.show_variables: 

203 print("# Resolved Variables:", file=sys.stderr) 

204 for key, value in sorted(cfg.variables.items()): 

205 print(f"# {key}: {value}", file=sys.stderr) 

206 print(file=sys.stderr) 

207 

208 # Format output 

209 if args.json: 

210 output_str = json.dumps(data, indent=2) 

211 else: 

212 output_str = yaml.dump( 

213 data, 

214 default_flow_style=False, 

215 sort_keys=False, 

216 allow_unicode=True, 

217 ) 

218 

219 # Write output 

220 if args.output: 

221 output_path = Path(args.output) 

222 output_path.write_text(output_str, encoding="utf-8") 

223 print(f"Output written to: {output_path}", file=sys.stderr) 

224 else: 

225 print(output_str) 

226 

227 return 0 

228 

229 except FileNotFoundError as e: 

230 print(f"Error: {e}", file=sys.stderr) 

231 return 1 

232 except ValueError as e: 

233 print(f"Error: {e}", file=sys.stderr) 

234 return 1 

235 except Exception as e: 

236 print(f"Unexpected error: {e}", file=sys.stderr) 

237 if "--debug" in (argv or sys.argv): 

238 raise 

239 return 1 

240 

241 

242if __name__ == "__main__": 

243 sys.exit(main()) 

244