diff --git a/flopy4/mf6/component.py b/flopy4/mf6/component.py index de3f0c29..d4b55972 100644 --- a/flopy4/mf6/component.py +++ b/flopy4/mf6/component.py @@ -1,10 +1,42 @@ from abc import ABC +from collections.abc import MutableMapping + +from xattree import xattree COMPONENTS = {} """MF6 component registry.""" -class Component(ABC): +@xattree +class Component(ABC, MutableMapping): @classmethod def __attrs_init_subclass__(cls): COMPONENTS[cls.__name__.lower()] = cls + + def __attrs_post_init__(self): + self._where = type(self).__xattree__["where"] + + def __getitem__(self, key): + data = getattr(self, self._where) + return data.children[key] + + def __setitem__(self, key, value): + data = getattr(self, self._where) + if key in data.children: + data.update({key: value}) + else: + data = data.assign({key: value}) + setattr(self, self._where, data) + + def __delitem__(self, key): + data = getattr(self, self._where) + data = data.drop_nodes(key) + setattr(self, self._where, data) + + def __iter__(self): + data = getattr(self, self._where) + return iter(data.children) + + def __len__(self): + data = getattr(self, self._where) + return len(data.children) diff --git a/test/test_component.py b/test/test_component.py index 3606878b..10edbe9c 100644 --- a/test/test_component.py +++ b/test/test_component.py @@ -234,3 +234,9 @@ def test_init_big_sim(): chd.head.data.todense(), sim.models["gwf"].chd[0].data.head.data.todense(), ) + + # test dictionary access/deletion + assert gwf.get("npf") + assert gwf["npf"].attrs["host"] is npf + del gwf["npf"] + assert not gwf.get("npf")