Coverage for ml_workbench/cli_datasets.py: 19%

1from __future__ import annotations 

2 

3import argparse 

4import json 

5from pathlib import Path 

6from typing import Any, Dict, Optional 

7 

8from .config import YamlConfig 

9from .dataset import Dataset 

10 

11 

12def _basic_stats_for_dataset(name: str, cfg: YamlConfig) -> dict[str, Any]: 

13 """Get basic statistics for a dataset using the Dataset class. 

14  

15 Parameters 

16 ---------- 

17 name : str 

18 Dataset name 

19 cfg : YamlConfig 

20 Configuration object 

21  

22 Returns 

23 ------- 

24 dict[str, Any] 

25 Dictionary with dataset statistics including name, description, format, 

26 type, path, num_columns, num_rows, column_names, and is_combined flag 

27 """ 

28 try: 

29 ds = Dataset(name, cfg) 

30 

31 result: dict[str, Any] = { 

32 "name": name, 

33 "description": ds.description, 

34 "format": ds.format, 

35 "type": ds.type, 

36 "path": ds.path, 

37 "is_combined": ds.is_combined, 

38 } 

39 

40 # Try to read and get statistics 

41 try: 

42 stats = ds.get_statistics() 

43 result.update({ 

44 "num_columns": stats["num_columns"], 

45 "num_rows": stats["num_rows"], 

46 "column_names": stats["column_names"], 

47 }) 

48 except Exception as exc: 

49 # If reading fails, still return basic info with error 

50 result.update({ 

51 "num_columns": None, 

52 "num_rows": None, 

53 "column_names": None, 

54 "error": str(exc), 

55 }) 

56 

57 return result 

58 

59 except Exception as exc: 

60 # If dataset creation fails, return minimal info 

61 return { 

62 "name": name, 

63 "error": f"Failed to create dataset: {exc}", 

64 } 

65 

66 

67def parse_kv_pairs(kv_list: list[str]) -> Dict[str, Any]: 

68 variables: Dict[str, Any] = {} 

69 for item in kv_list: 

70 if "=" not in item: 

71 raise argparse.ArgumentTypeError( 

72 f"Invalid variable format {item!r}, expected key=value" 

73 ) 

74 key, value = item.split("=", 1) 

75 variables[key] = value 

76 return variables 

77 

78 

79def build_arg_parser() -> argparse.ArgumentParser: 

80 parser = argparse.ArgumentParser( 

81 prog="ml-workbench-datasets", 

82 description="Inspect datasets from a YAML config and print basic statistics.", 

83 ) 

84 parser.add_argument("yaml", type=Path, help="Path to YAML configuration file") 

85 parser.add_argument( 

86 "--var", 

87 dest="vars", 

88 action="append", 

89 default=[], 

90 metavar="key=value", 

91 help="Placeholder variable override (can be used multiple times)", 

92 ) 

93 parser.add_argument( 

94 "--json", 

95 action="store_true", 

96 help="Output machine-readable JSON instead of pretty text", 

97 ) 

98 return parser 

99 

100 

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

102 parser = build_arg_parser() 

103 args = parser.parse_args(argv) 

104 

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

106 cfg = YamlConfig(args.yaml, **variables) 

107 

108 # Get list of dataset names and generate statistics for each 

109 dataset_names = cfg.get_datasets_list() 

110 results = [_basic_stats_for_dataset(name, cfg) for name in dataset_names] 

111 

112 if args.json: 

113 print(json.dumps({"datasets": results}, indent=2)) 

114 else: 

115 for item in results: 

116 print(f"- {item['name']}") 

117 if "error" in item: 

118 print(f" error: {item['error']}") 

119 continue 

120 

121 print(f" format: {item.get('format')}") 

122 print(f" type: {item.get('type')}") 

123 print(f" path: {item.get('path')}") 

124 

125 if item.get("is_combined"): 

126 print(" combined: true") 

127 

128 print(f" num_columns: {item.get('num_columns')}") 

129 print(f" num_rows: {item.get('num_rows')}") 

130 

131 # Print column names if available 

132 column_names = item.get("column_names") 

133 if column_names: 

134 print(f" columns: {', '.join(column_names)}") 

135 

136 return 0 

137 

138 

139if __name__ == "__main__": # pragma: no cover 

140 raise SystemExit(main()) 

141 

142