-
Notifications
You must be signed in to change notification settings - Fork 0
/
pdblist.py
35 lines (31 loc) · 1.01 KB
/
pdblist.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
from numpy import array
class PDBList(list):
"""
store as list of dict, but can function as a dict of list
"""
def __getitem__(self, key):
if type(key) is str:
#if key is str
try:
return [x[key] for x in self]
except:
return [[x[key] for x in y ] for y in self]
else:
try:
iter(key)
except TypeError:
#if key is not iterable or str
return super().__getitem__(key)
#if key is iterable but not str
return [[x[k] for k in key] for x in self]
def find_serial(self, serial):
for i, x in enumerate(self):
if x["serial"] == serial:
return i
def getposition(self, index):
entry = self[index]
return array((entry["x"], entry["y"], entry["z"]))
def setposition(self, index, value):
self[index]["x"] = value[0]
self[index]["y"] = value[1]
self[index]["z"] = value[2]