1from ase.io import read 

2import runpy 

3 

4 

5class CLICommand: 

6 """ Execute code on files. 

7 

8 The given python code is evaluated on the Atoms object read from 

9 the input file for each frame of the file. Either of -e or -E 

10 option should provided for evaluating code given as a string or 

11 from a file, respectively. 

12  

13 Variables which can be used inside the python code: 

14 - `index`: Index of the current Atoms object. 

15 - `atoms`: Current Atoms object. 

16 - `images`: List of all images given as input. 

17 """ 

18 

19 @staticmethod 

20 def add_arguments(parser): 

21 add = parser.add_argument 

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

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

24 help='Python code to execute on each atoms. The Atoms' 

25 ' object is available as `atoms`. ' 

26 'Example: For printing cell parameters from all the ' 

27 'frames, `print(atoms.cell.cellpar())`') 

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

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

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

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

32 help='Specify input FORMAT') 

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

34 default=':', metavar='NUMBER', 

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

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

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

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

39 '0:nimages:1.') 

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

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

42 help='Additional keyword arguments to pass to ' 

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

44 

45 @staticmethod 

46 def run(args, parser): 

47 if not (args.exec_code or args.exec_file): 

48 parser.error("At least one of '-e' or '-E' must be provided") 

49 

50 if args.read_args: 

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

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

53 

54 configs = [] 

55 for filename in args.input: 

56 atoms = read(filename, args.image_number, 

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

58 if isinstance(atoms, list): 

59 configs.extend(atoms) 

60 else: 

61 configs.append(atoms) 

62 

63 variables = {'images': configs} 

64 for index, atoms in enumerate(configs): 

65 variables['atoms'] = atoms 

66 variables['index'] = index 

67 if args.exec_code: 

68 exec(compile(args.exec_code, '<string>', 'exec'), variables) 

69 if args.exec_file: 

70 runpy.run_path(args.exec_file, init_globals=variables)