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
42