From 543eb99aa92fb1ff322ab6b4191f2a4c09236861 Mon Sep 17 00:00:00 2001 From: "David W.H. Swenson" Date: Thu, 31 Mar 2022 16:18:41 -0500 Subject: [PATCH 1/2] Implement __getitem__ for ChemicalSystem --- gufe/chemicalsystem.py | 3 +++ gufe/tests/test_chemicalsystem.py | 5 +++++ 2 files changed, 8 insertions(+) diff --git a/gufe/chemicalsystem.py b/gufe/chemicalsystem.py index a80ffb66..5d8f7402 100644 --- a/gufe/chemicalsystem.py +++ b/gufe/chemicalsystem.py @@ -127,6 +127,9 @@ def total_charge(self): total_charge += fc return total_charge + def __getitem__(self, item): + return self.components[item] + @classmethod def as_protein_smallmolecule_solvent(cls): """ """ diff --git a/gufe/tests/test_chemicalsystem.py b/gufe/tests/test_chemicalsystem.py index 290d1264..7442fae5 100644 --- a/gufe/tests/test_chemicalsystem.py +++ b/gufe/tests/test_chemicalsystem.py @@ -57,6 +57,8 @@ def test_ligand_construction(solv_comp, toluene_ligand_comp): 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): @@ -73,6 +75,9 @@ def test_complex_construction(prot_comp, solv_comp, toluene_ligand_comp): 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): From b4e093a05b8c7cf3fb7ad75c195691b19bf7bda8 Mon Sep 17 00:00:00 2001 From: "David W.H. Swenson" Date: Fri, 1 Apr 2022 09:50:02 -0500 Subject: [PATCH 2/2] make it a full mappings --- gufe/chemicalsystem.py | 9 ++++++++- gufe/tests/test_chemicalsystem.py | 6 ++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/gufe/chemicalsystem.py b/gufe/chemicalsystem.py index 5d8f7402..371fb8f8 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,12 @@ def total_charge(self): 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 7442fae5..4e89049c 100644 --- a/gufe/tests/test_chemicalsystem.py +++ b/gufe/tests/test_chemicalsystem.py @@ -54,6 +54,9 @@ 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 @@ -71,6 +74,9 @@ 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