From b2065add3c03e3e5a4328ec7f47cf001f283eedb Mon Sep 17 00:00:00 2001
From: Ilya Shlyakhter <ilya_shl@alum.mit.edu>
Date: Wed, 23 Aug 2017 16:17:34 -0400
Subject: [PATCH] added a more centralized way to preserve temp files for
 debugging (#659)

---
 util/cmd.py  |  3 ++-
 util/file.py | 36 ++++++++++++++++++++++++++++--------
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/util/cmd.py b/util/cmd.py
index 4b42a9dc8..bc3b59f2f 100644
--- a/util/cmd.py
+++ b/util/cmd.py
@@ -12,6 +12,7 @@
 import logging
 import argparse
 import util.version
+import util.file
 
 __author__ = "dpark@broadinstitute.org"
 __version__ = util.version.get_version()
@@ -200,7 +201,7 @@ def main_argparse(commands, description):
         try:
             ret = args.func_main(args)
         finally:
-            if hasattr(args, 'tmp_dirKeep') and args.tmp_dirKeep:
+            if (hasattr(args, 'tmp_dirKeep') and args.tmp_dirKeep) or util.file.keep_tmp():
                 log.exception(
                     "Exception occurred while running %s, saving tmp_dir at %s", args.command, tempfile.tempdir)
             else:
diff --git a/util/file.py b/util/file.py
index dbb45a36d..2627b21ab 100644
--- a/util/file.py
+++ b/util/file.py
@@ -82,7 +82,6 @@ def get_resources():
         resources = json.load(inf)
     return resources
 
-
 def mkstempfname(suffix='', prefix='tmp', directory=None, text=False):
     ''' There's no other one-liner way to securely ask for a temp file by
         filename only.  This calls mkstemp, which does what we want, except
@@ -111,9 +110,29 @@ def tempfnames(suffixes, *args, **kwargs):
     try:
         yield fns
     finally:
-        for fn in fns:
-            if os.path.isfile(fn):
-                os.unlink(fn)
+        if  not keep_tmp():
+            for fn in fns:
+                if os.path.isfile(fn):
+                    os.unlink(fn)
+
+@contextlib.contextmanager
+def tmp_dir(*args, **kwargs):
+    """Create and return a temporary directory, which is cleaned up on context exit
+    unless keep_tmp() is True."""
+    try:
+        name = tempfile.mkdtemp(*args, **kwargs)
+        yield name
+    finally:
+        if keep_tmp():
+            log.debug('keeping tempdir ' + name)
+        else:
+            shutil.rmtree(name)
+
+def keep_tmp():
+    """Whether to preserve temporary directories and files (useful during debugging).
+    Return True if the environment variable VIRAL_NGS_TMP_DIRKEEP is set.
+    """
+    return 'VIRAL_NGS_TMP_DIRKEEP' in os.environ
 
 def set_tmp_dir(name):
     proposed_prefix = ['tmp']
@@ -129,10 +148,11 @@ def set_tmp_dir(name):
 
 
 def destroy_tmp_dir(tempdir=None):
-    if tempdir:
-        shutil.rmtree(tempdir)
-    elif tempfile.tempdir:
-        shutil.rmtree(tempfile.tempdir)
+    if not keep_tmp():
+        if tempdir:
+            shutil.rmtree(tempdir)
+        elif tempfile.tempdir:
+            shutil.rmtree(tempfile.tempdir)
     tempfile.tempdir = None