1import os 

2 

3from ase.io import read, write 

4 

5 

6class CLICommand: 

7 """Convert between file formats. 

8 

9 Use "-" for stdin/stdout. 

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

11 """ 

12 

13 @staticmethod 

14 def add_arguments(parser): 

15 add = parser.add_argument 

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

17 help='Print names of converted files') 

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

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

20 help='Specify input FORMAT') 

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

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

23 help='Specify output FORMAT') 

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

25 help='Overwrite an existing file') 

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

27 default=':', metavar='NUMBER', 

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

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

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

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

32 '0:nimages:1.') 

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

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

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

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

37 'to suppress output of this frame.') 

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

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

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

41 add('-a', '--arrays', 

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

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

44 add('-I', '--info', 

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

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

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

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

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

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

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

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

53 help='Additional keyword arguments to pass to ' 

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

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

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

57 help='Additional keyword arguments to pass to ' 

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

59 

60 @staticmethod 

61 def run(args, parser): 

62 if args.verbose: 

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

64 if args.arrays: 

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

66 if args.verbose: 

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

68 if args.info: 

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

70 if args.verbose: 

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

72 if args.read_args: 

73 args.read_args = eval("dict({0})" 

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

75 if args.write_args: 

76 args.write_args = eval("dict({0})" 

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

78 

79 configs = [] 

80 for filename in args.input: 

81 atoms = read(filename, args.image_number, 

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

83 if isinstance(atoms, list): 

84 configs.extend(atoms) 

85 else: 

86 configs.append(atoms) 

87 

88 new_configs = [] 

89 for atoms in configs: 

90 if args.arrays: 

91 atoms.arrays = dict((k, atoms.arrays[k]) for k in args.arrays) 

92 if args.info: 

93 atoms.info = dict((k, atoms.info[k]) for k in args.info) 

94 if args.exec_code: 

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

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

97 if args.exec_file: 

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

99 'exec')) 

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

101 new_configs.append(atoms) 

102 configs = new_configs 

103 

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

105 parser.error('File already exists: {}'.format(args.output)) 

106 

107 if args.split_output: 

108 for i, atoms in enumerate(configs): 

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

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

111 else: 

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

113 **args.write_args)