1import os 

2import os.path as op 

3import sys 

4 

5from ase.io import read 

6from ase.io.formats import filetype, UnknownFileTypeError 

7from ase.db import connect 

8from ase.db.core import parse_selection 

9from ase.db.jsondb import JSONDatabase 

10from ase.db.row import atoms2dict 

11 

12 

13class CLICommand: 

14 """Find files with atoms in them. 

15 

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

17 

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

19 informations on how to construct the query string. 

20 """ 

21 

22 @staticmethod 

23 def add_arguments(parser): 

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

25 parser.add_argument( 

26 'query', nargs='?', 

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

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

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

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

31 help='More output.') 

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

33 help='Show also periodic boundary conditions, ' 

34 'chemical formula and filetype.') 

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

36 'ending with given strings. Example: ' 

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

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

39 'ending with given strings. Example: ' 

40 '"-x .cif".') 

41 

42 @staticmethod 

43 def run(args): 

44 main(args) 

45 

46 

47def main(args): 

48 query = parse_selection(args.query) 

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

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

51 

52 if args.long: 

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

54 

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

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

57 if format: 

58 if args.long: 

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

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

61 row.formula, 

62 format, 

63 path)) 

64 else: 

65 print(path) 

66 

67 

68def allpaths(folder, include, exclude): 

69 """Generate paths.""" 

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

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

72 for name in filenames: 

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

74 continue 

75 if include: 

76 for ext in include: 

77 if name.endswith(ext): 

78 break 

79 else: 

80 continue 

81 path = op.join(dirpath, name) 

82 yield path 

83 

84 # Skip .git, __pycache__ and friends: 

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

86 

87 

88def check(path, query, verbose): 

89 """Check a path. 

90 

91 Returns a (filetype, AtomsRow object) tuple. 

92 """ 

93 

94 try: 

95 format = filetype(path, guess=False) 

96 except (OSError, UnknownFileTypeError): 

97 return '', None 

98 

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

100 db = connect(path) 

101 else: 

102 try: 

103 atoms = read(path, format=format) 

104 except Exception as x: 

105 if verbose: 

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

107 return '', None 

108 db = FakeDB(atoms) 

109 

110 try: 

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

112 return format, row 

113 except Exception as x: 

114 if verbose: 

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

116 

117 return '', None 

118 

119 

120class FakeDB(JSONDatabase): 

121 def __init__(self, atoms): 

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

123 

124 def _read_json(self): 

125 return self.bigdct, [1], 2