Coverage for /builds/ase/ase/ase/io/db.py: 86.67%

1import ase.db 

2from ase.io.formats import string2index 

3 

4 

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

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

7 if isinstance(index, str): 

8 try: 

9 index = string2index(index) 

10 except ValueError: 

11 pass 

12 

13 if isinstance(index, int): 

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

15 

16 if isinstance(index, str): 

17 # index is a database query string: 

18 for row in db.select(index): 

19 yield row.toatoms() 

20 else: 

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

22 if start == stop: 

23 return 

24 assert step == 1 

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

26 yield row.toatoms() 

27 

28 

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

30 with ase.db.connect(filename, serial=True, append=append, **kwargs) as con: 

31 for atoms in images: 

32 con.write(atoms) 

33 

34 

35read_json = read_db 

36write_json = write_db 

37read_postgresql = read_db 

38write_postgresql = write_db 

39read_mysql = read_db 

40write_mysql = write_db