Repository: haddocking/arctic3d

Issue:
add run directory as input argument
this will be useful when excluding stuff. keep the default behaviour when no run_dir is selected by user

Discussion:

Patch:
diff --git a/src/arctic3d/cli.py b/src/arctic3d/cli.py
index ea51e99..cb6cded 100644
--- a/src/arctic3d/cli.py
+++ b/src/arctic3d/cli.py
@@ -56,6 +56,11 @@ argument_parser.add_argument(
     help="",
 )
 
+argument_parser.add_argument(
+    "--run_dir",
+    help="directory where to store the run",
+)
+
 
 def load_args(arguments):
     """
@@ -96,7 +101,7 @@ def maincli():
     cli(argument_parser, main)
 
 
-def main(input_arg, db, interface_file, out_uniprot, out_pdb, pdb_to_use):
+def main(input_arg, db, interface_file, out_uniprot, out_pdb, pdb_to_use, run_dir):
     """Main function."""
     log.setLevel("DEBUG")
 
@@ -119,7 +124,7 @@ def main(input_arg, db, interface_file, out_uniprot, out_pdb, pdb_to_use):
 
     log.info(f"Target UNIPROTID: {uniprot_id}")
 
-    input_files = setup_output_folder(uniprot_id, input_files)
+    input_files = setup_output_folder(uniprot_id, input_files, run_dir)
 
     # retrieve interfaces
     if interface_file:
diff --git a/src/arctic3d/modules/output.py b/src/arctic3d/modules/output.py
index 33209ef..6a280e4 100644
--- a/src/arctic3d/modules/output.py
+++ b/src/arctic3d/modules/output.py
@@ -11,7 +11,7 @@ import plotly.graph_objects as go
 log = logging.getLogger("arctic3dlog")
 
 
-def setup_output_folder(uniprot_id, input_files):
+def setup_output_folder(uniprot_id, input_files, output_dir):
     """Sets up output folder.
 
     Parameters
@@ -20,19 +20,27 @@ def setup_output_folder(uniprot_id, input_files):
         uniprot_id of the run
     input_files : dict of Paths
         dict of input files
+    output_dir : str or None
+        user-defined name of the run
     Returns
     -------
     copied_input_files : dict of Paths
         dict of copied input files
     """
-    if uniprot_id is None:
-        key = "custom"
-    else:
-        key = uniprot_id
-    run_dir = Path(f"arctic3d-{key}")
+    run_dir = output_dir
+    if run_dir is None:
+        if uniprot_id is None:
+            key = "custom"
+        else:
+            key = uniprot_id
+        run_dir = f"arctic3d-{key}"
+    run_dir = Path(run_dir)
+
     if os.path.exists(run_dir):
         raise Exception(f"{run_dir} already exists!")
+
     # setting up the directory
+    log.info(f"Setting up output_directory {run_dir}")
     os.mkdir(run_dir)
     datadir = Path(run_dir, "input_data")
     os.mkdir(datadir)
@@ -41,6 +49,7 @@ def setup_output_folder(uniprot_id, input_files):
         filepath = Path(datadir, filename)
         shutil.copy(input_files[key], filepath)
     os.chdir(run_dir)
+
     # retrieving info about copied input files
     copied_input_files = {}
     for key in input_files:


Test patch:
diff --git a/tests/test_output.py b/tests/test_output.py
index eb7a8be..a0b6ce1 100644
--- a/tests/test_output.py
+++ b/tests/test_output.py
@@ -5,6 +5,7 @@ import pytest
 
 from arctic3d.modules.output import (
     output_pdb,
+    setup_output_folder,
     write_clusters,
     write_residues,
     write_residues_probs,
@@ -94,3 +95,17 @@ def test_output_pdb(inp_pdb):
     for ln_id in range(len(observed_content)):
         assert original_content[ln_id][:60] == observed_content[ln_id][:60]
     os.unlink(output_files[0])
+
+
+def test_run_dir():
+    """Test if the expected run_dir is effectively created."""
+    run_dir = "run_dir"
+    uniprot_id = "fake_uniprot"
+    start_cwd = os.getcwd()
+    setup_output_folder(uniprot_id, [], run_dir)
+    obs_cwd = Path(os.getcwd())
+    exp_cwd = Path(start_cwd, run_dir)
+    assert exp_cwd == obs_cwd
+    os.chdir(start_cwd)
+    os.rmdir(Path(run_dir, "input_data"))
+    os.rmdir(run_dir)