Coverage for /builds/ase/ase/ase/cli/exec.py: 93.75%

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 """ Execute code on files. 

10 

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

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

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

14 from a file, respectively. 

15 

16 Variables which can be used inside the python code: 

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

18 - `atoms`: Current Atoms object. 

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

20 """ 

21 

22 @staticmethod 

23 def add_arguments(parser): 

24 add = parser.add_argument 

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

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

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

28 ' object is available as `atoms`. ' 

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

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

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

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

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

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

35 help='Specify input FORMAT') 

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

37 default=':', metavar='NUMBER', 

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

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

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

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

42 '0:nimages:1.') 

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

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

45 help='Additional keyword arguments to pass to ' 

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

47 

48 @staticmethod 

49 def run(args, parser): 

50 import runpy 

51 

52 from ase.io import read 

53 

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

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

56 

57 if args.read_args: 

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

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

60 

61 configs = [] 

62 for filename in args.input: 

63 atoms = read(filename, args.image_number, 

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

65 if isinstance(atoms, list): 

66 configs.extend(atoms) 

67 else: 

68 configs.append(atoms) 

69 

70 variables = {'images': configs} 

71 for index, atoms in enumerate(configs): 

72 variables['atoms'] = atoms 

73 variables['index'] = index 

74 if args.exec_code: 

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

76 if args.exec_file: 

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