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

1# Note: 

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

3# import time during CLI execution 

4import sys 

5 

6 

7class CLICommand: 

8 """Find files with atoms in them. 

9 

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

11 

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

13 informations on how to construct the query string. 

14 """ 

15 

16 @staticmethod 

17 def add_arguments(parser): 

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

19 parser.add_argument( 

20 'query', nargs='?', 

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

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

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

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

25 help='More output.') 

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

27 help='Show also periodic boundary conditions, ' 

28 'chemical formula and filetype.') 

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

30 'ending with given strings. Example: ' 

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

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

33 'ending with given strings. Example: ' 

34 '"-x .cif".') 

35 

36 @staticmethod 

37 def run(args): 

38 main(args) 

39 

40 

41def main(args): 

42 from ase.db.core import parse_selection 

43 

44 query = parse_selection(args.query) 

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

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

47 

48 if args.long: 

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

50 

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

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

53 if format: 

54 if args.long: 

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

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

57 row.formula, 

58 format, 

59 path)) 

60 else: 

61 print(path) 

62 

63 

64def allpaths(folder, include, exclude): 

65 """Generate paths.""" 

66 import os 

67 import os.path as op 

68 

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

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

71 for name in filenames: 

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

73 continue 

74 if include: 

75 for ext in include: 

76 if name.endswith(ext): 

77 break 

78 else: 

79 continue 

80 path = op.join(dirpath, name) 

81 yield path 

82 

83 # Skip .git, __pycache__ and friends: 

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

85 

86 

87def check(path, query, verbose): 

88 """Check a path. 

89 

90 Returns a (filetype, AtomsRow object) tuple. 

91 """ 

92 from ase.db import connect 

93 from ase.db.jsondb import JSONDatabase 

94 from ase.db.row import atoms2dict 

95 from ase.io import read 

96 from ase.io.formats import UnknownFileTypeError, filetype 

97 

98 class FakeDB(JSONDatabase): 

99 def __init__(self, atoms): 

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

101 

102 def _read_json(self): 

103 return self.bigdct, [1], 2 

104 

105 try: 

106 format = filetype(path, guess=False) 

107 except (OSError, UnknownFileTypeError): 

108 return '', None 

109 

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

111 db = connect(path) 

112 else: 

113 try: 

114 atoms = read(path, format=format) 

115 except Exception as x: 

116 if verbose: 

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

118 return '', None 

119 db = FakeDB(atoms) 

120 

121 try: 

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

123 return format, row 

124 except Exception as x: 

125 if verbose: 

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

127 

128 return '', None