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

1# Note: 

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

3# import time during CLI execution 

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 import os 

63 

64 from ase.io import read, write 

65 

66 if args.verbose: 

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

68 if args.arrays: 

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

70 if args.verbose: 

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

72 if args.info: 

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

74 if args.verbose: 

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

76 if args.read_args: 

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

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

79 if args.write_args: 

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

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

82 

83 configs = [] 

84 for filename in args.input: 

85 atoms = read(filename, args.image_number, 

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

87 if isinstance(atoms, list): 

88 configs.extend(atoms) 

89 else: 

90 configs.append(atoms) 

91 

92 new_configs = [] 

93 for atoms in configs: 

94 if args.arrays: 

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

96 if args.info: 

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

98 if args.exec_code: 

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

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

101 if args.exec_file: 

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

103 'exec')) 

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

105 new_configs.append(atoms) 

106 configs = new_configs 

107 

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

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

110 

111 if args.split_output: 

112 for i, atoms in enumerate(configs): 

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

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

115 else: 

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

117 **args.write_args)