diff --git a/gufe/chemicalsystem.py b/gufe/chemicalsystem.py index 68563c81..ce408a72 100644 --- a/gufe/chemicalsystem.py +++ b/gufe/chemicalsystem.py @@ -1,6 +1,7 @@ # This code is part of OpenFE and is licensed under the MIT license. # For details, see https://github.com/OpenFreeEnergy/gufe +from collections import abc from typing import Dict, Optional import numpy as np @@ -9,7 +10,7 @@ from .component import Component -class ChemicalSystem(Serializable): +class ChemicalSystem(Serializable, abc.Mapping): """A node of an alchemical network. Attributes @@ -130,6 +131,15 @@ def total_charge(self): total_charge += fc return total_charge + def __getitem__(self, item): + return self.components[item] + + def __iter__(self): + return iter(self.components) + + def __len__(self): + return len(self.components) + @classmethod def as_protein_smallmolecule_solvent(cls): """ """ diff --git a/gufe/tests/test_chemicalsystem.py b/gufe/tests/test_chemicalsystem.py index 240c6b5b..486b187a 100644 --- a/gufe/tests/test_chemicalsystem.py +++ b/gufe/tests/test_chemicalsystem.py @@ -54,9 +54,14 @@ def test_ligand_construction(solv_comp, toluene_ligand_comp): ) assert len(state.components) == 2 + assert len(state) == 2 + + assert list(state) == ['solvent', 'ligand'] assert state.components['solvent'] == solv_comp assert state.components['ligand'] == toluene_ligand_comp + assert state['solvent'] == solv_comp + assert state['ligand'] == toluene_ligand_comp def test_complex_construction(prot_comp, solv_comp, toluene_ligand_comp): @@ -69,10 +74,16 @@ def test_complex_construction(prot_comp, solv_comp, toluene_ligand_comp): ) assert len(state.components) == 3 + assert len(state) == 3 + + assert list(state) == ['protein', 'solvent', 'ligand'] assert state.components['protein'] == prot_comp assert state.components['solvent'] == solv_comp assert state.components['ligand'] == toluene_ligand_comp + assert state['protein'] == prot_comp + assert state['solvent'] == solv_comp + assert state['ligand'] == toluene_ligand_comp def test_hash_and_eq(prot_comp, solv_comp, toluene_ligand_comp):