Merge pull request #15 from OpenFreeEnergy/chemsys-getitem

Make `ChemicalSystem` an `abc.Mapping`

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):
         """ """

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):