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

1# Note: 

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

3# import time during CLI execution 

4 

5 

6class CLICommand: 

7 """ Execute code on files. 

8 

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

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

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

12 from a file, respectively. 

13 

14 Variables which can be used inside the python code: 

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

16 - `atoms`: Current Atoms object. 

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

18 """ 

19 

20 @staticmethod 

21 def add_arguments(parser): 

22 add = parser.add_argument 

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

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

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

26 ' object is available as `atoms`. ' 

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

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

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

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

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

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

33 help='Specify input FORMAT') 

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

35 default=':', metavar='NUMBER', 

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

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

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

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

40 '0:nimages:1.') 

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

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

43 help='Additional keyword arguments to pass to ' 

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

45 

46 @staticmethod 

47 def run(args, parser): 

48 import runpy 

49 

50 from ase.io import read 

51 

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

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

54 

55 if args.read_args: 

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

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

58 

59 configs = [] 

60 for filename in args.input: 

61 atoms = read(filename, args.image_number, 

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

63 if isinstance(atoms, list): 

64 configs.extend(atoms) 

65 else: 

66 configs.append(atoms) 

67 

68 variables = {'images': configs} 

69 for index, atoms in enumerate(configs): 

70 variables['atoms'] = atoms 

71 variables['index'] = index 

72 if args.exec_code: 

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

74 if args.exec_file: 

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