diff --git a/src/openfe/storage/warehouse.py b/src/openfe/storage/warehouse.py index 02a81d5c4..c10baa12c 100644 --- a/src/openfe/storage/warehouse.py +++ b/src/openfe/storage/warehouse.py @@ -25,6 +25,8 @@ class WarehouseStores(TypedDict): ---------- setup : ExternalStorage Storage location for setup-related objects and configurations. + result : ExternalStorage + Storage location for result-related object. Notes ----- @@ -32,7 +34,7 @@ class WarehouseStores(TypedDict): """ setup: ExternalStorage - # We will add a result and task store here in the future. + result: ExternalStorage class WarehouseBaseClass: @@ -63,7 +65,7 @@ def __repr__(self): # probably should include repr of external store, too return f"{self.__class__.__name__}({self.stores})" - def delete(self, store_name: Literal["setup"], location: str): + def delete(self, store_name: Literal["setup", "result"], location: str): """Delete an object from a specific store. Parameters @@ -106,6 +108,31 @@ def load_setup_tokenizable(self, obj: GufeKey) -> GufeTokenizable: """ return self._load_gufe_tokenizable(gufe_key=obj) + def store_result_tokenizable(self, obj: GufeTokenizable): + """Store a GufeTokenizable object from the result store. + + Parameters + ---------- + obj : GufeKey + The key of the object to store. + """ + return self._store_gufe_tokenizable("result", obj) + + def load_result_tokenizable(self, obj: GufeKey) -> GufeTokenizable: + """Load a GufeTokenizable object from the result store. + + Parameters + ---------- + obj : GufeKey + The key of the object to load. + + Returns + ------- + GufeTokenizable + The loaded object. + """ + return self._load_gufe_tokenizable(gufe_key=obj) + def exists(self, key: GufeKey) -> bool: """Check if an object with the given key exists in any store. @@ -144,15 +171,15 @@ def _get_store_for_key(self, key: GufeKey) -> ExternalStorage: return self.stores[name] raise ValueError(f"GufeKey {key} is not stored") - def _store_gufe_tokenizable(self, store_name: Literal["setup"], obj: GufeTokenizable): + def _store_gufe_tokenizable(self, store_name: Literal["setup", "result"], obj: GufeTokenizable): """Store a GufeTokenizable object with deduplication. - Parameters - ---------- - store_name : Literal["setup"] - Name of the store to store the object in. - obj : GufeTokenizable - The object to store. + Parameters + ---------- + store_name : Literal["setup"] + Name of the store to store the object in. + obj : GufeTokenizable + The object to store. Notes ----- @@ -246,15 +273,26 @@ def recursive_build_object_cache(key: GufeKey) -> GufeTokenizable: @property def setup_store(self): - """Get the setup store. + """Get the setup store Returns ------- ExternalStorage - The setup storage location. + The setup storage location """ return self.stores["setup"] + @property + def result_store(self): + """Get the result store. + + Returns + ------- + ExternalStorage + The result storage location + """ + return self.stores["result"] + class FileSystemWarehouse(WarehouseBaseClass): """Warehouse implementation using local filesystem storage. @@ -276,7 +314,6 @@ class FileSystemWarehouse(WarehouseBaseClass): def __init__(self, root_dir: str = "warehouse"): setup_store = FileStorage(f"{root_dir}/setup") - # When we add a result store it will look like this - # result_store = FileStorage(f"{root_dir}/results") - stores = WarehouseStores(setup=setup_store) + result_store = FileStorage(f"{root_dir}/result") + stores = WarehouseStores(setup=setup_store, result=result_store) super().__init__(stores)