Coverage for aipyapp/cli/command/base_parser.py: 0%

1import argparse 

2from collections import OrderedDict 

3 

4from .base import BaseCommand, Completable 

5 

6def requires_value(action): 

7 ret = ( 

8 action.nargs is not None or 

9 (action.type is not None and action.type != bool) or 

10 action.choices is not None 

11 ) 

12 return ret 

13 

14class ParserCommand(BaseCommand): 

15 """Base class for all commands""" 

16 def __init__(self, manager=None): 

17 super().__init__(manager) 

18 self.parser = None 

19 self.arguments = None 

20 self.subcommands = None 

21 

22 def get_subcommands(self): 

23 """Get subcommands""" 

24 return self.subcommands 

25 

26 def init(self): 

27 """Initialize the command, can be overridden by subclasses""" 

28 parser = argparse.ArgumentParser(prog=f'/{self.name}', description=self.description, exit_on_error=False) 

29 self.add_arguments(parser) 

30 if hasattr(self, 'add_subcommands'): 

31 subparsers = parser.add_subparsers(dest='subcommand') 

32 self.add_subcommands(subparsers) 

33 

34 arguments = OrderedDict() 

35 for action in parser._actions: 

36 # 处理选项参数（如 --option） 

37 for option in action.option_strings: 

38 if option in ('-h', '--help'): 

39 continue 

40 

41 choices = OrderedDict() 

42 if action.choices: 

43 for choice in action.choices: 

44 choices[choice] = Completable(choice) 

45 arguments[option] = Completable(option, action.help, choices=choices, requires_value=requires_value(action)) 

46 

47 # 处理位置参数（如 'name'） 

48 if not action.option_strings and action.dest != 'help' and action.dest != 'subcommand': 

49 choices = OrderedDict() 

50 if action.choices: 

51 for choice in action.choices: 

52 choices[choice] = Completable(choice) 

53 # 对于位置参数，不将参数名作为自动补齐的一部分，直接使用选项值 

54 if action.choices: 

55 for choice in action.choices: 

56 arguments[choice] = Completable(choice, action.help, requires_value=False) 

57 else: 

58 arguments[action.dest] = Completable(action.dest, action.help, choices=choices, requires_value=requires_value(action)) 

59 

60 subcommands = OrderedDict() 

61 for action in parser._actions: 

62 if not isinstance(action, argparse._SubParsersAction): 

63 continue 

64 

65 for subaction in action._get_subactions(): 

66 cmd_name = subaction.dest or subaction.name 

67 subcommands[cmd_name] = Completable(cmd_name, subaction.help) 

68 

69 for subcmd, subparser in action.choices.items(): 

70 sub_arguments = OrderedDict() 

71 for sub_action in subparser._actions: 

72 # 处理子命令的选项参数 

73 for option in sub_action.option_strings: 

74 if option in ('-h', '--help'): 

75 continue 

76 

77 choices = OrderedDict() 

78 if sub_action.choices: 

79 for choice in sub_action.choices: 

80 choices[choice] = Completable(choice) 

81 

82 sub_arguments[option] = Completable(option, sub_action.help, choices=choices, requires_value=requires_value(sub_action)) 

83 

84 # 处理子命令的位置参数 

85 if not sub_action.option_strings and sub_action.dest != 'help': 

86 choices = OrderedDict() 

87 if sub_action.choices: 

88 for choice in sub_action.choices: 

89 choices[choice] = Completable(choice) 

90 # 对于位置参数，不将参数名作为自动补齐的一部分，直接使用选项值 

91 if sub_action.choices: 

92 for choice in sub_action.choices: 

93 sub_arguments[choice] = Completable(choice, sub_action.help, requires_value=False) 

94 else: 

95 sub_arguments[sub_action.dest] = Completable(sub_action.dest, sub_action.help, choices=choices, requires_value=requires_value(sub_action)) 

96 

97 # 将 arguments 作为属性添加到 Completable 对象中 

98 if subcmd in subcommands: 

99 subcommands[subcmd]['arguments'] = sub_arguments 

100 

101 self.parser = parser 

102 self.arguments = arguments 

103 self.subcommands = subcommands 

104 

105 def add_arguments(self, parser): 

106 """Add command-specific arguments to the parser""" 

107 pass 

108 

109 def get_arg_values(self, arg, subcommand=None, partial_value=''): 

110 """Get argument values for argument `arg`""" 

111 choices = arg.get('choices') 

112 if choices: 

113 return choices.values() 

114 return None 

115 

116 def execute(self, args): 

117 """Execute the command with parsed arguments""" 

118 subcommand = getattr(args, 'subcommand', None) 

119 if subcommand: 

120 func = getattr(self, f'cmd_{subcommand}', None) 

121 if not func: 

122 self.log.error(f"Subcommand {subcommand} not found") 

123 return 

124 else: 

125 func = self.cmd 

126 

127 return func(args, ctx=self.manager.context) 

128 

129 def cmd(self, args, ctx): 

130 """Execute the main command""" 

131 pass