[WIP: don't review] New storage in network planner CLI

openfecli/commands/plan_rbfe_network.py

@@ -7,6 +7,7 @@
 from openfecli import OFECommandPlugin
 from openfecli.parameters import (
     MOL_DIR, PROTEIN, MAPPER, OUTPUT_DIR, COFACTORS,
+    NEW_STORAGE_OUTPUT,  # separate line for easy delete later
 )
 from openfecli.plan_alchemical_networks_utils import plan_alchemical_network_output
 
@@ -82,9 +83,11 @@
     help=OUTPUT_DIR.kwargs["help"] + " Defaults to `./alchemicalNetwork`.",
     default="alchemicalNetwork",
 )
+@NEW_STORAGE_OUTPUT.parameter()
 @print_duration
 def plan_rbfe_network(
-    molecules: List[str], protein: str, cofactors: tuple[str], output_dir: str
+    molecules: List[str], protein: str, cofactors: tuple[str],
+    output_dir: str, new_storage: bool,
 ):
     """
     Plan a relative binding free energy network, saved as JSON files for
@@ -119,6 +122,8 @@
     from openfe.setup.atom_mapping.lomap_scorers import (
         default_lomap_score,
     )
+    from openfe.storage.resultclient import ResultClient
+    from gufe.storage.externalresource import FileStorage
     from openfe.setup import LomapAtomMapper
     from openfe.setup.ligand_network_planning import (
         generate_minimal_spanning_network,
@@ -175,12 +180,18 @@
 
     # OUTPUT
     write("Output:")
-    write("\tSaving to: " + str(output_dir))
-    plan_alchemical_network_output(
-        alchemical_network=alchemical_network,
-        ligand_network=ligand_network,
-        folder_path=OUTPUT_DIR.get(output_dir),
-    )
+    if new_storage:
+        write(f"Saving to: {output_dir}")
+        storage = ResultClient(FileStorage(output_dir))
+        storage.store_network(alchemical_network)
+        # TODO: save out the ligand network as well
+    else:
+        write("\tSaving to: " + str(output_dir))
+        plan_alchemical_network_output(
+            alchemical_network=alchemical_network,
+            ligand_network=ligand_network,
+            folder_path=OUTPUT_DIR.get(output_dir),
+        )
 
 
 PLUGIN = OFECommandPlugin(

openfecli/parameters/__init__.py

@@ -4,6 +4,6 @@
 from .mol import MOL
 from .mapper import MAPPER
 from .output import OUTPUT_FILE_AND_EXT
-from .output_dir import OUTPUT_DIR
+from .output_dir import OUTPUT_DIR, NEW_STORAGE_OUTPUT
 from .protein import PROTEIN
 from .molecules import MOL_DIR, COFACTORS

openfecli/parameters/output_dir.py

@@ -16,3 +16,10 @@ def get_dir(user_input, context):
     getter=get_dir,
     type=click.Path(file_okay=False, resolve_path=True),
 )
+
+NEW_STORAGE_OUTPUT = Option(
+    "--new-storage",
+    help="use the new storage",
+    is_flag=True,
+    hidden=True,
+)