Skip to content

Commit 79c55cf

Browse files
authored
dictlike behavior for components (#130)
basics for #117. maybe later we can use the dict proxy class in xattree instead of implementing it ourselves?
1 parent a5e7b8f commit 79c55cf

File tree

2 files changed

+39
-1
lines changed

2 files changed

+39
-1
lines changed

flopy4/mf6/component.py

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,42 @@
11
from abc import ABC
2+
from collections.abc import MutableMapping
3+
4+
from xattree import xattree
25

36
COMPONENTS = {}
47
"""MF6 component registry."""
58

69

7-
class Component(ABC):
10+
@xattree
11+
class Component(ABC, MutableMapping):
812
@classmethod
913
def __attrs_init_subclass__(cls):
1014
COMPONENTS[cls.__name__.lower()] = cls
15+
16+
def __attrs_post_init__(self):
17+
self._where = type(self).__xattree__["where"]
18+
19+
def __getitem__(self, key):
20+
data = getattr(self, self._where)
21+
return data.children[key]
22+
23+
def __setitem__(self, key, value):
24+
data = getattr(self, self._where)
25+
if key in data.children:
26+
data.update({key: value})
27+
else:
28+
data = data.assign({key: value})
29+
setattr(self, self._where, data)
30+
31+
def __delitem__(self, key):
32+
data = getattr(self, self._where)
33+
data = data.drop_nodes(key)
34+
setattr(self, self._where, data)
35+
36+
def __iter__(self):
37+
data = getattr(self, self._where)
38+
return iter(data.children)
39+
40+
def __len__(self):
41+
data = getattr(self, self._where)
42+
return len(data.children)

test/test_component.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -234,3 +234,9 @@ def test_init_big_sim():
234234
chd.head.data.todense(),
235235
sim.models["gwf"].chd[0].data.head.data.todense(),
236236
)
237+
238+
# test dictionary access/deletion
239+
assert gwf.get("npf")
240+
assert gwf["npf"].attrs["host"] is npf
241+
del gwf["npf"]
242+
assert not gwf.get("npf")

0 commit comments

Comments
 (0)