Coverage for /builds/ase/ase/ase/cli/find.py: 20.55%

1# fmt: off 

2 

3# Note: 

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

5# import time during CLI execution 

6import sys 

7 

8 

9class CLICommand: 

10 """Find files with atoms in them. 

11 

12 Search through files known to ASE applying a query to filter the results. 

13 

14 See https://wiki.fysik.dtu.dk/ase/ase/db/db.html#querying for more 

15 informations on how to construct the query string. 

16 """ 

17 

18 @staticmethod 

19 def add_arguments(parser): 

20 parser.add_argument('folder', help='Folder to look in.') 

21 parser.add_argument( 

22 'query', nargs='?', 

23 help='Examples: More than 2 hydrogens and no silver: "H>2,Ag=0". ' 

24 'More than 1000 atoms: "natoms>1000". ' 

25 'Slab geometry containing Cu and Ni: "pbc=TTF,Cu,Ni".') 

26 parser.add_argument('-v', '--verbose', action='store_true', 

27 help='More output.') 

28 parser.add_argument('-l', '--long', action='store_true', 

29 help='Show also periodic boundary conditions, ' 

30 'chemical formula and filetype.') 

31 parser.add_argument('-i', '--include', help='Include only filenames ' 

32 'ending with given strings. Example: ' 

33 '"-i .xyz,.traj".') 

34 parser.add_argument('-x', '--exclude', help='Exclude filenames ' 

35 'ending with given strings. Example: ' 

36 '"-x .cif".') 

37 

38 @staticmethod 

39 def run(args): 

40 main(args) 

41 

42 

43def main(args): 

44 from ase.db.core import parse_selection 

45 

46 query = parse_selection(args.query) 

47 include = args.include.split(',') if args.include else [] 

48 exclude = args.exclude.split(',') if args.exclude else [] 

49 

50 if args.long: 

51 print('pbc {:10} {:15} path'.format('formula', 'filetype')) 

52 

53 for path in allpaths(args.folder, include, exclude): 

54 format, row = check(path, query, args.verbose) 

55 if format: 

56 if args.long: 

57 print('{} {:10} {:15} {}' 

58 .format(''.join(str(p) for p in row.pbc.astype(int)), 

59 row.formula, 

60 format, 

61 path)) 

62 else: 

63 print(path) 

64 

65 

66def allpaths(folder, include, exclude): 

67 """Generate paths.""" 

68 import os 

69 import os.path as op 

70 

71 exclude += ['.py', '.pyc'] 

72 for dirpath, dirnames, filenames in os.walk(folder): 

73 for name in filenames: 

74 if any(name.endswith(ext) for ext in exclude): 

75 continue 

76 if include: 

77 for ext in include: 

78 if name.endswith(ext): 

79 break 

80 else: 

81 continue 

82 path = op.join(dirpath, name) 

83 yield path 

84 

85 # Skip .git, __pycache__ and friends: 

86 dirnames[:] = (name for name in dirnames if name[0] not in '._') 

87 

88 

89def check(path, query, verbose): 

90 """Check a path. 

91 

92 Returns a (filetype, AtomsRow object) tuple. 

93 """ 

94 from ase.db import connect 

95 from ase.db.jsondb import JSONDatabase 

96 from ase.db.row import atoms2dict 

97 from ase.io import read 

98 from ase.io.formats import UnknownFileTypeError, filetype 

99 

100 class FakeDB(JSONDatabase): 

101 def __init__(self, atoms): 

102 self.bigdct = {1: atoms2dict(atoms)} 

103 

104 def _read_json(self): 

105 return self.bigdct, [1], 2 

106 

107 try: 

108 format = filetype(path, guess=False) 

109 except (OSError, UnknownFileTypeError): 

110 return '', None 

111 

112 if format in ['db', 'json']: 

113 db = connect(path) 

114 else: 

115 try: 

116 atoms = read(path, format=format) 

117 except Exception as x: 

118 if verbose: 

119 print(path + ':', x, file=sys.stderr) 

120 return '', None 

121 db = FakeDB(atoms) 

122 

123 try: 

124 for row in db._select(*query): 

125 return format, row 

126 except Exception as x: 

127 if verbose: 

128 print(path + ':', x, file=sys.stderr) 

129 

130 return '', None