def _select(self, keys, cmps, explain=False, verbosity=0,
limit=None, offset=0, sort=None):
if explain:
yield {'explain': (0, 0, 0, 'scan table')}
return
if sort:
if sort[0] == '-':
reverse = True
sort = sort[1:]
else:
reverse = False
def f(row):
return row[sort]
rows = sorted(self._select(keys + [sort], cmps),
key=f, reverse=reverse)
if limit:
rows = rows[offset:offset + limit]
for row in rows:
yield row
return
try:
bigdct, ids, nextid = self._read_json()
except IOError:
return
if not limit:
limit = -offset - 1
cmps = [(key, ops[op], val) for key, op, val in cmps]
n = 0
for id in ids:
if n - offset == limit:
return
row = AtomsRow(bigdct[id])
row.id = id
for key in keys:
if key not in row:
break
else:
for key, op, val in cmps:
if isinstance(key, int):
value = np.equal(row.numbers, key).sum()
else:
value = row.get(key)
if key == 'pbc':
assert op in [ops['='], ops['!=']]
value = ''.join('FT'[x] for x in value)
if value is None or not op(value, val):
break
else:
if n >= offset:
yield row
n += 1
def _select(self,
keys,
cmps,
explain=False,
verbosity=0,
limit=None,
offset=0,
sort=None,
include_data=True):
if explain:
yield {'explain': (0, 0, 0, 'scan table')}
return
if sort:
if sort[0] == '-':
reverse = True
sort = sort[1:]
else:
reverse = False
def f(row):
return row[sort]
rows = sorted(self._select(keys + [sort], cmps),
key=f,
reverse=reverse)
if limit:
rows = rows[offset:offset + limit]
for row in rows:
yield row
return
try:
bigdct, ids, nextid = self._read_json()
except IOError:
return
if not limit:
limit = -offset - 1
cmps = [(key, ops[op], val) for key, op, val in cmps]
n = 0
for id in ids:
if n - offset == limit:
return
dct = bigdct[id]
if not include_data:
dct.pop('data', None)
row = AtomsRow(dct)
row.id = id
for key in keys:
if key not in row:
break
else:
for key, op, val in cmps:
if isinstance(key, int):
value = np.equal(row.numbers, key).sum()
else:
value = row.get(key)
if key == 'pbc':
assert op in [ops['='], ops['!=']]
value = ''.join('FT'[x] for x in value)
if value is None or not op(value, val):
break
else:
if n >= offset:
yield row
n += 1
def update(self,
id,
atoms=None,
delete_keys=[],
data=None,
**add_key_value_pairs):
"""Update and/or delete key-value pairs of row(s).
id: int
ID of row to update.
atoms: Atoms object
Optionally update the Atoms data (positions, cell, ...).
data: dict
Data dict to be added to the existing data.
delete_keys: list of str
Keys to remove.
Use keyword arguments to add new key-value pairs.
Returns number of key-value pairs added and removed.
"""
if not isinstance(id, numbers.Integral):
if isinstance(id, list):
err = ('First argument must be an int and not a list.\n'
'Do something like this instead:\n\n'
'with db:\n'
' for id in ids:\n'
' db.update(id, ...)')
raise ValueError(err)
raise TypeError('id must be an int')
check(add_key_value_pairs)
row = self._get_row(id)
kvp = row.key_value_pairs
n = len(kvp)
for key in delete_keys:
kvp.pop(key, None)
n -= len(kvp)
m = -len(kvp)
kvp.update(add_key_value_pairs)
m += len(kvp)
moredata = data
data = row.get('data', {})
if moredata:
data.update(moredata)
if not data:
data = None
if atoms:
oldrow = row
row = AtomsRow(atoms)
# Copy over data, kvp, ctime, user and id
row._data = oldrow._data
row.__dict__.update(kvp)
row._keys = list(kvp)
row.ctime = oldrow.ctime
row.user = oldrow.user
row.id = id
if atoms or os.path.splitext(self.filename)[1] == '.json':
self._write(row, kvp, data, row.id)
else:
self._update(row.id, kvp, data)
return m, n
