db/message: make sure not to leave temporary files around

Use a contextlib.ExitStack to remove the temporary file used by external
handler, even if an exception is raised.

1 files changed, 32 insertions, 34 deletions

diff --git a/alot/db/message.py b/alot/db/message.py
index 33b0f880..ff8cca70 100644
--- a/alot/db/message.py
+++ b/alot/db/message.py
@@ -1,6 +1,8 @@
 # Copyright (C) 2011-2012  Patrick Totzke <patricktotzke@gmail.com>
 # This file is released under the GNU GPL, version 3 or a later revision.
 # For further details see the COPYING file
+
+from   contextlib import ExitStack
 import email
 import email.charset as charset
 import email.policy
@@ -39,42 +41,38 @@ def _render_part_external(raw_payload, ctype, params, field_key='copiousoutput')
     if isinstance(raw_payload, str):
         raw_payload = raw_payload.encode('utf-8')
 
-    tempfile_name = None
-    stdin = None
-    handler_raw_commandstring = entry['view']
-    # in case the mailcap defined command contains no '%s',
-    # we pipe the files content to the handling command via stdin
-    if '%s' in handler_raw_commandstring:
-        # open tempfile, respect mailcaps nametemplate
-        nametemplate = entry.get('nametemplate', '%s')
-        prefix, suffix = parse_mailcap_nametemplate(nametemplate)
-        with tempfile.NamedTemporaryFile(
-                delete=False, prefix=prefix, suffix=suffix) \
-                as tmpfile:
+    with ExitStack() as stack:
+        # read parameter, create handler command
+        parms = tuple('='.join(p) for p in params)
+
+        # in case the mailcap defined command contains no '%s',
+        # we pipe the files content to the handling command via stdin
+        if '%s' in entry['view']:
+            # open tempfile, respect mailcaps nametemplate
+            nametemplate   = entry.get('nametemplate', '%s')
+            prefix, suffix = parse_mailcap_nametemplate(nametemplate)
+
+            tmpfile = stack.enter_context(tempfile.NamedTemporaryFile(prefix = prefix, suffix = suffix))
+
             tmpfile.write(raw_payload)
+            tmpfile.flush()
+
             tempfile_name = tmpfile.name
-    else:
-        stdin = raw_payload
-
-    # read parameter, create handler command
-    parms = tuple('='.join(p) for p in params)
-
-    # create and call external command
-    cmd = mailcap.subst(entry['view'], ctype,
-                        filename=tempfile_name, plist=parms)
-    logging.debug('command: %s', cmd)
-    logging.debug('parms: %s', str(parms))
-    cmdlist = split_commandstring(cmd)
-    # call handler
-    stdout, _, _ = helper.call_cmd(cmdlist, stdin=stdin)
-    if stdout:
-        rendered_payload = stdout
-
-    # remove tempfile
-    if tempfile_name:
-        os.unlink(tempfile_name)
-
-    return rendered_payload
+            stdin         = None
+        else:
+            tempfile_name = None
+            stdin         = raw_payload
+
+        # create and call external command
+        cmd = mailcap.subst(entry['view'], ctype,
+                            filename = tempfile_name, plist = parms)
+        logging.debug('command: %s', cmd)
+        logging.debug('parms: %s', str(parms))
+        cmdlist = split_commandstring(cmd)
+        # call handler
+        stdout, _, _ = helper.call_cmd(cmdlist, stdin=stdin)
+
+        return stdout
 
 class _MessageHeaders:
     _msg = None