From a3e37fd4ac822d5edda9cda6cec3e9306fbfef63 Mon Sep 17 00:00:00 2001
From: Will Dumm <wrhdumm@gmail.com>
Date: Tue, 27 Aug 2024 11:19:15 -0800
Subject: [PATCH] 131 Fix rendering hang when parent directory of output file
 doesn't exist (#132)

* create parent folders if necessary in render

* create nonexistant outbase parent directories if necessary

* fix tests
---
 gctree/branching_processes.py | 14 +++++++++-----
 gctree/cli.py                 |  4 ++++
 tests/smalltest.sh            |  2 +-
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/gctree/branching_processes.py b/gctree/branching_processes.py
index b382b90..e2956ff 100755
--- a/gctree/branching_processes.py
+++ b/gctree/branching_processes.py
@@ -18,7 +18,6 @@
 import numpy as np
 import warnings
 import random
-import os
 import scipy.special as scs
 import scipy.optimize as sco
 import ete3
@@ -36,6 +35,7 @@
 from typing import Tuple, Dict, List, Union, Set, Callable, Mapping, Sequence, Optional
 from decimal import Decimal
 import math
+from pathlib import Path
 
 sequence_resolutions = hdag.parsimony_utils.standard_nt_ambiguity_map_gap_as_char.get_sequence_resolution_func(
     "sequence"
@@ -485,6 +485,10 @@ def render(
             show_nuc_muts: If True, annotate branches with nucleotide mutations. If False, and frame is provided, then branches
                 will be annotated with amino acid mutations.
         """
+        # Create parent directories of outfile if they don't exist
+        outfile = Path(outfile)
+        outfile.parent.mkdir(parents=True, exist_ok=True)
+
         if frame is not None and frame not in (1, 2, 3):
             raise RuntimeError("frame must be 1, 2, or 3")
 
@@ -675,10 +679,10 @@ def my_layout(node):
                         description=f"abundance={node.abundance}",
                     )
                 )
-            AlignIO.write(
-                aln, open(os.path.splitext(outfile)[0] + ".fasta", "w"), "fasta"
-            )
-        return tree_copy.render(outfile, tree_style=ts)
+            with open(outfile.with_suffix(".fasta"), "w") as fh:
+                AlignIO.write(aln, fh, "fasta")
+
+        return tree_copy.render(str(outfile), tree_style=ts)
 
     def feature_colormap(
         self,
diff --git a/gctree/cli.py b/gctree/cli.py
index 904ef95..f1e5b44 100644
--- a/gctree/cli.py
+++ b/gctree/cli.py
@@ -15,6 +15,7 @@
 import random
 import ete3
 import itertools
+from pathlib import Path
 
 
 def test(args):
@@ -164,6 +165,9 @@ def isotype_add(forest):
             isotype_names=args.isotype_names,
         )
 
+    # Make parent directory of output file base string if it doesn't exist
+    Path(args.outbase).parent.mkdir(parents=True, exist_ok=True)
+
     if len(args.infiles) == 2:
         forest = bp.CollapsedForest(
             pp.parse_outfile(args.infiles[0], args.infiles[1], args.root)
diff --git a/tests/smalltest.sh b/tests/smalltest.sh
index 9011f4a..5b22aaa 100755
--- a/tests/smalltest.sh
+++ b/tests/smalltest.sh
@@ -16,7 +16,7 @@ wget -O $substitutions https://bitbucket.org/kleinstein/shazam/raw/ba4b30fc6791e
 
 gctree infer tests/small_outfile tests/abundances.csv --outbase tests/smalltest_output/gctree.infer --root GL --frame 1 --verbose --idlabel --idmapfile tests/idmap.txt --isotype_mapfile tests/isotypemap.txt --mutability HS5F_Mutability.csv --substitution HS5F_Substitution.csv --ranking_coeffs 1 1 0 --use_old_mut_parsimony --branching_process_ranking_coeff 0 
 
-gctree infer tests/smalltest_output/gctree.infer.inference.parsimony_forest.p --outbase tests/smalltest_output/gctree.infer --root GL --frame 1 --verbose --idlabel --idmapfile tests/idmap.txt --isotype_mapfile tests/isotypemap.txt --mutability HS5F_Mutability.csv --substitution HS5F_Substitution.csv --ranking_coeffs 1 -1 0 --use_old_mut_parsimony
+gctree infer tests/smalltest_output/gctree.infer.inference.parsimony_forest.p --outbase tests/smalltest_output/nonexistant_dir/gctree.infer --root GL --frame 1 --verbose --idlabel --idmapfile tests/idmap.txt --isotype_mapfile tests/isotypemap.txt --mutability HS5F_Mutability.csv --substitution HS5F_Substitution.csv --ranking_coeffs 1 -1 0 --use_old_mut_parsimony
 
 gctree infer tests/small_outfile tests/abundances.csv --outbase tests/smalltest_output/gctree.infer --root GL --frame 1 --verbose --idlabel --idmapfile tests/idmap.txt --isotype_mapfile tests/isotypemap.txt --mutability HS5F_Mutability.csv --substitution HS5F_Substitution.csv --use_old_mut_parsimony