Coverage for /builds/kinetik161/ase/ase/cli/main.py: 65.22%

69 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-12-10 11:04 +0000

1import argparse 

2import textwrap 

3from importlib import import_module 

4 

5from ase import __version__ 

6 

7 

8class CLIError(Exception): 

9 """Error for CLI commands. 

10 

11 A subcommand may raise this. The message will be forwarded to 

12 the error() method of the argument parser.""" 

13 

14 

15# Important: Following any change to command-line parameters, use 

16# python3 -m ase.cli.completion to update autocompletion. 

17commands = [ 

18 ('info', 'ase.cli.info'), 

19 # ('show', 'ase.cli.show'), 

20 ('test', 'ase.test'), 

21 ('gui', 'ase.gui.ag'), 

22 ('db', 'ase.cli.db'), 

23 ('run', 'ase.cli.run'), 

24 ('band-structure', 'ase.cli.band_structure'), 

25 ('build', 'ase.cli.build'), 

26 ('dimensionality', 'ase.cli.dimensionality'), 

27 ('eos', 'ase.eos'), 

28 ('ulm', 'ase.cli.ulm'), 

29 ('find', 'ase.cli.find'), 

30 ('nebplot', 'ase.cli.nebplot'), 

31 ('nomad-upload', 'ase.cli.nomad'), 

32 ('nomad-get', 'ase.cli.nomadget'), 

33 ('convert', 'ase.cli.convert'), 

34 ('reciprocal', 'ase.cli.reciprocal'), 

35 ('completion', 'ase.cli.completion'), 

36 ('diff', 'ase.cli.diff'), 

37 ('exec', 'ase.cli.exec') 

38] 

39 

40 

41def main(prog='ase', description='ASE command line tool.', 

42 version=__version__, commands=commands, hook=None, args=None): 

43 parser = argparse.ArgumentParser(prog=prog, 

44 description=description, 

45 formatter_class=Formatter) 

46 parser.add_argument('--version', action='version', 

47 version=f'%(prog)s-{version}') 

48 parser.add_argument('-T', '--traceback', action='store_true') 

49 subparsers = parser.add_subparsers(title='Sub-commands', 

50 dest='command') 

51 

52 subparser = subparsers.add_parser('help', 

53 description='Help', 

54 help='Help for sub-command.') 

55 subparser.add_argument('helpcommand', 

56 nargs='?', 

57 metavar='sub-command', 

58 help='Provide help for sub-command.') 

59 

60 functions = {} 

61 parsers = {} 

62 for command, module_name in commands: 

63 cmd = import_module(module_name).CLICommand 

64 docstring = cmd.__doc__ 

65 if docstring is None: 

66 # Backwards compatibility with GPAW 

67 short = cmd.short_description 

68 long = getattr(cmd, 'description', short) 

69 else: 

70 parts = docstring.split('\n', 1) 

71 if len(parts) == 1: 

72 short = docstring 

73 long = docstring 

74 else: 

75 short, body = parts 

76 long = short + '\n' + textwrap.dedent(body) 

77 subparser = subparsers.add_parser( 

78 command, 

79 formatter_class=Formatter, 

80 help=short, 

81 description=long) 

82 cmd.add_arguments(subparser) 

83 functions[command] = cmd.run 

84 parsers[command] = subparser 

85 

86 if hook: 

87 args = hook(parser, args) 

88 else: 

89 args = parser.parse_args(args) 

90 

91 if args.command == 'help': 

92 if args.helpcommand is None: 

93 parser.print_help() 

94 else: 

95 parsers[args.helpcommand].print_help() 

96 elif args.command is None: 

97 parser.print_usage() 

98 else: 

99 f = functions[args.command] 

100 try: 

101 if f.__code__.co_argcount == 1: 

102 f(args) 

103 else: 

104 f(args, parsers[args.command]) 

105 except KeyboardInterrupt: 

106 pass 

107 except CLIError as x: 

108 parser.error(x) 

109 except Exception as x: 

110 if args.traceback: 

111 raise 

112 else: 

113 l1 = f'{x.__class__.__name__}: {x}\n' 

114 l2 = ('To get a full traceback, use: {} -T {} ...' 

115 .format(prog, args.command)) 

116 parser.error(l1 + l2) 

117 

118 

119class Formatter(argparse.HelpFormatter): 

120 """Improved help formatter.""" 

121 

122 def _fill_text(self, text, width, indent): 

123 assert indent == '' 

124 out = '' 

125 blocks = text.split('\n\n') 

126 for block in blocks: 

127 if block[0] == '*': 

128 # List items: 

129 for item in block[2:].split('\n* '): 

130 out += textwrap.fill(item, 

131 width=width - 2, 

132 initial_indent='* ', 

133 subsequent_indent=' ') + '\n' 

134 elif block[0] == ' ': 

135 # Indented literal block: 

136 out += block + '\n' 

137 else: 

138 # Block of text: 

139 out += textwrap.fill(block, width=width) + '\n' 

140 out += '\n' 

141 return out[:-1]