Merge pull request #362 from OpenFreeEnergy/add-keyedchain

Added some convenience methods to KeyedChain for symmetry; test simplifications

gufe/tests/test_tokenization.py

@@ -401,9 +401,7 @@ def test_gufe_objects_from_shallow_dict(solvated_complex):
     gufe_objects = set(gufe_objects_from_shallow_dict(shallow_dict))
 
     assert len(gufe_objects) == 3
-
-    for gufe_object in gufe_objects:
-        assert gufe_object in solvated_complex.components.values()
+    assert set(gufe_objects) == set(solvated_complex.components.values())
 
 
 class TestKeyedChain:
@@ -422,11 +420,13 @@ def test_from_gufe(self, benzene_variants_star_map):
         ]
 
         kc_gufe_keys = set(kc.gufe_keys())
-        kc_shallow_dicts = list(kc.keyed_dicts())
+        kc_keyed_dicts = list(kc.keyed_dicts())
+
+        assert kc_gufe_keys == set(original_keys)
 
         for key, keyed_dict in zip(original_keys, original_keyed_dicts):
             assert key in kc_gufe_keys
-            assert keyed_dict in kc_shallow_dicts
+            assert keyed_dict in kc_keyed_dicts
 
     def test_to_gufe(self, benzene_variants_star_map):
         kc = KeyedChain.from_gufe(benzene_variants_star_map)

gufe/tokenization.py

@@ -721,7 +721,7 @@ def add_edges(o):
 
 
 class KeyedChain(object):
-    """Keyed chain representation of a GufeTokenizable.
+    """Keyed chain representation encoder of a GufeTokenizable.
 
     The keyed chain representation of a GufeTokenizable provides a
     topologically sorted list of gufe keys and GufeTokenizable keyed dicts
@@ -731,6 +731,31 @@ class KeyedChain(object):
     The class wraps around a list of tuples containing the gufe key and the
     keyed dict form of the GufeTokenizable.
 
+    Examples
+    --------
+    We can create a keyed chain representation from any GufeTokenizable, such
+    as:
+
+    >>> from gufe.tokenization import KeyedChain
+    >>> s = SolventComponent()
+    >>> keyed_chain = KeyedChain.gufe_to_keyed_chain_rep(s)
+    >>> keyed_chain
+    [('SolventComponent-26b4034ad9dbd9f908dfc298ea8d449f',
+      {'smiles': 'O',
+       'positive_ion': 'Na+',
+       'negative_ion': 'Cl-',
+       'ion_concentration': '0.15 molar',
+       'neutralize': True,
+       '__qualname__': 'SolventComponent',
+       '__module__': 'gufe.components.solventcomponent',
+       ':version:': 1})]
+
+    And we can do the reverse operation as well to go from a keyed chain
+    representation back to a GufeTokenizable:
+
+    >>> KeyedChain(keyed_chain).to_gufe()
+    SolventComponent(name=O, Na+, Cl-)
+
     """
 
     def __init__(self, keyed_chain):
@@ -749,6 +774,15 @@ def to_gufe(self) -> GufeTokenizable:
             gts[gufe_key] = gt
         return gt
 
+    @classmethod
+    def from_keyed_chain_rep(cls, keyed_chain: List[Tuple[str, Dict]]) -> Self:
+        """Initialize a KeyedChain from a keyed chain representation."""
+        return cls(keyed_chain)
+
+    def to_keyed_chain_rep(self) -> List[Tuple[str, Dict]]:
+        """Return the keyed chain representation of this object."""
+        return list(self)
+
     @staticmethod
     def gufe_to_keyed_chain_rep(
         gufe_object: GufeTokenizable,