diff --git a/alchemiscale/base/api.py b/alchemiscale/base/api.py index a44865e3..b8e49417 100644 --- a/alchemiscale/base/api.py +++ b/alchemiscale/base/api.py @@ -14,7 +14,7 @@ from fastapi import status as http_status from fastapi.routing import APIRoute from fastapi.security import OAuth2PasswordRequestForm -from gufe.tokenization import JSON_HANDLER, GufeTokenizable +from gufe.tokenization import JSON_HANDLER, GufeTokenizable, KeyedChain from ..settings import ( JWTSettings, @@ -32,7 +32,6 @@ oauth2_scheme, ) from ..security.models import Token, TokenData, CredentialedEntity -from ..keyedchain import KeyedChain def validate_scopes(scope: Scope, token: TokenData) -> None: diff --git a/alchemiscale/interface/api.py b/alchemiscale/interface/api.py index fef0b2c3..5b6aeb1e 100644 --- a/alchemiscale/interface/api.py +++ b/alchemiscale/interface/api.py @@ -12,7 +12,7 @@ from fastapi.middleware.gzip import GZipMiddleware import json -from gufe.tokenization import JSON_HANDLER +from gufe.tokenization import JSON_HANDLER, KeyedChain from ..base.api import ( GufeJSONResponse, @@ -34,7 +34,6 @@ from ..storage.models import TaskStatusEnum from ..models import Scope, ScopedKey from ..security.models import TokenData, CredentialedUserIdentity -from ..keyedchain import KeyedChain app = FastAPI(title="AlchemiscaleAPI") diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index eaba27b7..7bd1311f 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -13,7 +13,7 @@ from async_lru import alru_cache import networkx as nx from gufe import AlchemicalNetwork, Transformation, ChemicalSystem -from gufe.tokenization import GufeTokenizable, JSON_HANDLER +from gufe.tokenization import GufeTokenizable, JSON_HANDLER, KeyedChain from gufe.protocols import ProtocolResult, ProtocolDAGResult @@ -29,7 +29,6 @@ ) from ..strategies import Strategy from ..validators import validate_network_nonself -from ..keyedchain import KeyedChain from warnings import warn diff --git a/alchemiscale/keyedchain.py b/alchemiscale/keyedchain.py deleted file mode 100644 index e2cf187c..00000000 --- a/alchemiscale/keyedchain.py +++ /dev/null @@ -1,82 +0,0 @@ -from gufe.tokenization import GufeTokenizable, key_decode_dependencies -import networkx as nx -from alchemiscale.utils import gufe_to_digraph - -from typing import List, Tuple, Dict, Generator - - -class KeyedChain(object): - """Keyed chain representation of a GufeTokenizable. - - The keyed chain representation of a GufeTokenizable provides a - topologically sorted list of gufe keys and GufeTokenizable keyed dicts - that can be used to fully recreate a GufeTokenizable without the need for a - populated TOKENIZATION_REGISTRY. - - The class wraps around a list of tuples containing the gufe key and the - keyed dict form of the GufeTokenizable. - - """ - - def __init__(self, keyed_chain): - self._keyed_chain = keyed_chain - - @classmethod - def from_gufe(cls, gufe_object: GufeTokenizable) -> super: - """Initialize a KeyedChain from a GufeTokenizable.""" - return cls(cls.gufe_to_keyed_chain_rep(gufe_object)) - - def to_gufe(self) -> GufeTokenizable: - """Initialize a GufeTokenizable.""" - gts = {} - for gufe_key, keyed_dict in self: - gt = key_decode_dependencies(keyed_dict, registry=gts) - gts[gufe_key] = gt - return gt - - @staticmethod - def gufe_to_keyed_chain_rep( - gufe_object: GufeTokenizable, - ) -> List[Tuple[str, Dict]]: - """Create the keyed chain represenation of a GufeTokenizable. - - This represents the GufeTokenizable as a list of two-element tuples - containing, as their first and second elements, the gufe key and keyed - dict form of the GufeTokenizable, respectively, and provides the - underlying structure used in the KeyedChain class. - - Parameters - ---------- - gufe_object - The GufeTokenizable for which the KeyedChain is generated. - - Returns - ------- - key_and_keyed_dicts - The keyed chain represenation of a GufeTokenizable. - - """ - key_and_keyed_dicts = [ - (str(gt.key), gt.to_keyed_dict()) - for gt in nx.topological_sort(gufe_to_digraph(gufe_object)) - ][::-1] - return key_and_keyed_dicts - - def gufe_keys(self) -> Generator[str, None, None]: - """Create a generator that iterates over the gufe keys in the KeyedChain.""" - for key, _ in self: - yield key - - def keyed_dicts(self) -> Generator[Dict, None, None]: - """Create a generator that iterates over the keyed dicts in the KeyedChain.""" - for _, _dict in self: - yield _dict - - def __len__(self): - return len(self._keyed_chain) - - def __iter__(self): - return self._keyed_chain.__iter__() - - def __getitem__(self, index): - return self._keyed_chain[index] diff --git a/alchemiscale/tests/integration/interface/test_api.py b/alchemiscale/tests/integration/interface/test_api.py index 7814b8c7..6702a99e 100644 --- a/alchemiscale/tests/integration/interface/test_api.py +++ b/alchemiscale/tests/integration/interface/test_api.py @@ -2,10 +2,9 @@ import json from gufe import AlchemicalNetwork, ChemicalSystem, Transformation -from gufe.tokenization import JSON_HANDLER, GufeTokenizable +from gufe.tokenization import JSON_HANDLER, GufeTokenizable, KeyedChain from alchemiscale.models import ScopedKey -from alchemiscale.keyedchain import KeyedChain def pre_load_payload(network, scope, name="incomplete 2"): diff --git a/alchemiscale/tests/unit/test_keyedchain.py b/alchemiscale/tests/unit/test_keyedchain.py deleted file mode 100644 index d272d5b6..00000000 --- a/alchemiscale/tests/unit/test_keyedchain.py +++ /dev/null @@ -1,45 +0,0 @@ -from alchemiscale.keyedchain import KeyedChain -from alchemiscale.utils import RegistryBackup - -from gufe.tokenization import get_all_gufe_objs, GufeTokenizable, is_gufe_key_dict - - -def test_keyedchain_full_network(network): - objects = get_all_gufe_objs(network) - - for o in objects: - with RegistryBackup(): - kc = KeyedChain.from_gufe(o) - _o = kc.to_gufe() - assert o == _o - - -def test_keyedchain_len(network): - objects = get_all_gufe_objs(network) - expect_len = len(objects) - - keyedchain = KeyedChain.from_gufe(network) - - assert len(keyedchain) == expect_len - - -def test_keyedchain_get_keys(network): - keyedchain = KeyedChain.from_gufe(network) - keys = list(map(lambda x: x.key, get_all_gufe_objs(network))) - - assert set(keyedchain.gufe_keys()) == set(keys) - - -def test_keyedchain_get_keyed_dicts(network): - keyedchain = KeyedChain.from_gufe(network) - - for keyed_dict in keyedchain.keyed_dicts(): - assert isinstance(keyed_dict, dict) - - -def test_keyedchain_iteration(network): - keyedchain = KeyedChain.from_gufe(network) - - for key, keyed_dict in keyedchain: - gt = GufeTokenizable.from_keyed_dict(keyed_dict) - assert gt.key == key