1import ase.db 

2from ase.io.formats import string2index 

3 

4 

5def read_db(filename, index, **kwargs): 

6 db = ase.db.connect(filename, serial=True, **kwargs) 

7 

8 if isinstance(index, str): 

9 try: 

10 index = string2index(index) 

11 except ValueError: 

12 pass 

13 

14 if isinstance(index, int): 

15 index = slice(index, index + 1 or None) 

16 

17 if isinstance(index, str): 

18 # index is a database query string: 

19 for row in db.select(index): 

20 yield row.toatoms() 

21 else: 

22 start, stop, step = index.indices(db.count()) 

23 if start == stop: 

24 return 

25 assert step == 1 

26 for row in db.select(offset=start, limit=stop - start): 

27 yield row.toatoms() 

28 

29 

30def write_db(filename, images, append=False, **kwargs): 

31 con = ase.db.connect(filename, serial=True, append=append, **kwargs) 

32 for atoms in images: 

33 con.write(atoms) 

34 

35 

36read_json = read_db 

37write_json = write_db 

38read_postgresql = read_db 

39write_postgresql = write_db 

40read_mysql = read_db 

41write_mysql = write_db