Coverage for /builds/ase/ase/ase/cli/convert.py: 72.13%

1# fmt: off 

2 

3# Note: 

4# Try to avoid module level import statements here to reduce 

5# import time during CLI execution 

6 

7 

8class CLICommand: 

9 """Convert between file formats. 

10 

11 Use "-" for stdin/stdout. 

12 See "ase info --formats" for known formats. 

13 """ 

14 

15 @staticmethod 

16 def add_arguments(parser): 

17 add = parser.add_argument 

18 add('-v', '--verbose', action='store_true', 

19 help='Print names of converted files') 

20 add('input', nargs='+', metavar='input-file') 

21 add('-i', '--input-format', metavar='FORMAT', 

22 help='Specify input FORMAT') 

23 add('output', metavar='output-file') 

24 add('-o', '--output-format', metavar='FORMAT', 

25 help='Specify output FORMAT') 

26 add('-f', '--force', action='store_true', 

27 help='Overwrite an existing file') 

28 add('-n', '--image-number', 

29 default=':', metavar='NUMBER', 

30 help='Pick images from trajectory. NUMBER can be a ' 

31 'single number (use a negative number to count from ' 

32 'the back) or a range: start:stop:step, where the ' 

33 '":step" part can be left out - default values are ' 

34 '0:nimages:1.') 

35 add('-e', '--exec-code', 

36 help='Python code to execute on each atoms before ' 

37 'writing it to output file. The Atoms object is ' 

38 'available as `atoms`. Set `atoms.info["_output"] = False` ' 

39 'to suppress output of this frame.') 

40 add('-E', '--exec-file', 

41 help='Python source code file to execute on each ' 

42 'frame, usage is as for -e/--exec-code.') 

43 add('-a', '--arrays', 

44 help='Comma-separated list of atoms.arrays entries to include ' 

45 'in output file. Default is all entries.') 

46 add('-I', '--info', 

47 help='Comma-separated list of atoms.info entries to include ' 

48 'in output file. Default is all entries.') 

49 add('-s', '--split-output', action='store_true', 

50 help='Write output frames to individual files. ' 

51 'Output file name should be a format string with ' 

52 'a single integer field, e.g. out-{:0>5}.xyz') 

53 add('--read-args', nargs='+', action='store', 

54 default={}, metavar="KEY=VALUE", 

55 help='Additional keyword arguments to pass to ' 

56 '`ase.io.read()`.') 

57 add('--write-args', nargs='+', action='store', 

58 default={}, metavar="KEY=VALUE", 

59 help='Additional keyword arguments to pass to ' 

60 '`ase.io.write()`.') 

61 

62 @staticmethod 

63 def run(args, parser): 

64 import os 

65 

66 from ase.io import read, write 

67 

68 if args.verbose: 

69 print(', '.join(args.input), '->', args.output) 

70 if args.arrays: 

71 args.arrays = [k.strip() for k in args.arrays.split(',')] 

72 if args.verbose: 

73 print('Filtering to include arrays: ', ', '.join(args.arrays)) 

74 if args.info: 

75 args.info = [k.strip() for k in args.info.split(',')] 

76 if args.verbose: 

77 print('Filtering to include info: ', ', '.join(args.info)) 

78 if args.read_args: 

79 args.read_args = eval("dict({})" 

80 .format(', '.join(args.read_args))) 

81 if args.write_args: 

82 args.write_args = eval("dict({})" 

83 .format(', '.join(args.write_args))) 

84 

85 configs = [] 

86 for filename in args.input: 

87 atoms = read(filename, args.image_number, 

88 format=args.input_format, **args.read_args) 

89 if isinstance(atoms, list): 

90 configs.extend(atoms) 

91 else: 

92 configs.append(atoms) 

93 

94 new_configs = [] 

95 for atoms in configs: 

96 if args.arrays: 

97 atoms.arrays = {k: atoms.arrays[k] for k in args.arrays} 

98 if args.info: 

99 atoms.info = {k: atoms.info[k] for k in args.info} 

100 if args.exec_code: 

101 # avoid exec() for Py 2+3 compat. 

102 eval(compile(args.exec_code, '<string>', 'exec')) 

103 if args.exec_file: 

104 eval(compile(open(args.exec_file).read(), args.exec_file, 

105 'exec')) 

106 if "_output" not in atoms.info or atoms.info["_output"]: 

107 new_configs.append(atoms) 

108 configs = new_configs 

109 

110 if not args.force and os.path.isfile(args.output): 

111 parser.error(f'File already exists: {args.output}') 

112 

113 if args.split_output: 

114 for i, atoms in enumerate(configs): 

115 write(args.output.format(i), atoms, 

116 format=args.output_format, **args.write_args) 

117 else: 

118 write(args.output, configs, format=args.output_format, 

119 **args.write_args)