Coverage for src/utils/cli.py: 96%

1# Copyright (c) CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL 

2# Unauthorized copying of this file and/or project, via any medium is strictly prohibited. 

3# Direct all questions to legal@cloudzero.com 

4 

5""" 

6Command line interface utilities for CloudZero Data Tool. 

7Handles argument parsing and validation. 

8""" 

9 

10import argparse 

11import sys 

12from pathlib import Path 

13from rich.console import Console 

14 

15console = Console() 

16 

17 

18def create_argument_parser() -> argparse.ArgumentParser: 

19 """Create and configure the argument parser for czdt.""" 

20 parser = argparse.ArgumentParser( 

21 description="Analyze Perplexity CSV data with interactive column exploration" 

22 ) 

23 parser.add_argument( 

24 "--in", 

25 type=Path, 

26 dest="input_file", 

27 help="Path to CSV file to analyze (required for analysis operations)", 

28 ) 

29 parser.add_argument( 

30 "--column", type=str, help="Specific column to analyze (skips interactive mode)" 

31 ) 

32 parser.add_argument( 

33 "--analysis", 

34 action="store_true", 

35 help="Run interactive CloudZero telemetry mapping analysis to identify how CSV data can be mapped to CloudZero telemetry records", 

36 ) 

37 parser.add_argument( 

38 "--costformation", 

39 type=Path, 

40 help="Optional costformation.yaml file to suggest cost dimensions for telemetry mapping", 

41 ) 

42 parser.add_argument( 

43 "--config", 

44 type=Path, 

45 help="Configuration file (JSON/YAML) containing predefined answers for telemetry analysis automation", 

46 ) 

47 parser.add_argument( 

48 "--rollup", 

49 type=str, 

50 help="Create rollup/pivot table. Format: 'group_cols:agg_specs' (e.g., 'MODEL_SKU,RUNTIME_MODEL_NAME:count(LLM_INPUT_TOKENS),sum(CT_TOTAL_COMPLETION_TOKENS)'). Use with --group-by for time-based analysis.", 

51 ) 

52 parser.add_argument( 

53 "--rollup-limit", 

54 type=int, 

55 default=None, 

56 help="Maximum number of rollup rows to display (default: show all)", 

57 ) 

58 parser.add_argument( 

59 "--out", 

60 type=str, 

61 help="Export rollup results to CSV file instead of displaying on screen", 

62 ) 

63 parser.add_argument( 

64 "--group-by", 

65 choices=["hour", "day"], 

66 help="Group rollup data by time period (hour or day). Requires --rollup option.", 

67 ) 

68 parser.add_argument( 

69 "--datetime-col", 

70 type=str, 

71 help="Specify datetime column for time grouping. If not provided, will auto-detect.", 

72 ) 

73 parser.add_argument( 

74 "--allow-none", 

75 action="store_true", 

76 help="Allow rows with None/null values in non-aggregate columns when creating rollup. By default, null values are ignored.", 

77 ) 

78 parser.add_argument( 

79 "--lazy", 

80 action="store_true", 

81 help="Use lazy loading for better memory efficiency with very large datasets", 

82 ) 

83 parser.add_argument( 

84 "--fast", 

85 action="store_true", 

86 help="Enable all performance optimizations (lazy loading, column selection, etc.)", 

87 ) 

88 parser.add_argument( 

89 "--sample", 

90 type=int, 

91 help="Process only a sample of N rows for faster analysis of large files", 

92 ) 

93 parser.add_argument( 

94 "--generate-schema", 

95 type=Path, 

96 help="Generate schema manifest for all CSV files in the specified folder", 

97 ) 

98 parser.add_argument( 

99 "--interactive", 

100 action="store_true", 

101 help="Start interactive mode for column exploration", 

102 ) 

103 

104 return parser 

105 

106 

107def validate_arguments(args) -> bool: 

108 """Validate parsed arguments and show appropriate errors.""" 

109 # Validate group-by requires rollup 

110 if args.group_by and not args.rollup: 

111 console.print( 

112 "❌ [red]Error: --group-by option requires --rollup to be specified[/red]" 

113 ) 

114 return False 

115 

116 # Validate allow-none requires rollup 

117 if args.allow_none and not args.rollup: 

118 console.print( 

119 "❌ [red]Error: --allow-none option requires --rollup to be specified[/red]" 

120 ) 

121 return False 

122 

123 # Check if analysis operation requires --in argument 

124 analysis_operations = [args.analysis, args.rollup, args.column, args.interactive] 

125 if any(analysis_operations) and not args.input_file: 

126 console.print( 

127 "❌ [red]Error: --in argument is required for analysis operations[/red]" 

128 ) 

129 return False 

130 

131 return True 

132 

133 

134def should_show_help(args) -> bool: 

135 """Determine if help should be shown (no meaningful arguments provided).""" 

136 return not any( 

137 [ 

138 args.analysis, 

139 args.rollup, 

140 args.column, 

141 args.interactive, 

142 args.generate_schema, 

143 ] 

144 ) 

145 

146 

147def parse_and_validate_args(): 

148 """Parse command line arguments and validate them.""" 

149 parser = create_argument_parser() 

150 args = parser.parse_args() 

151 

152 # Handle schema generation first (doesn't require --in argument) 

153 if args.generate_schema: 

154 return args, True # Valid, special case 

155 

156 # Check if no arguments provided, show help and exit 

157 if should_show_help(args): 

158 parser.print_help() 

159 sys.exit(0) 

160 

161 # Validate arguments 

162 if not validate_arguments(args): 

163 sys.exit(1) 

164 

165 return args, True