From 309fb25e9b089618c37f1a741fa6009cce54ac9e Mon Sep 17 00:00:00 2001
From: Anton Khirnov <anton@khirnov.net>
Date: Sat, 30 Jan 2021 16:33:56 +0100
Subject: db/attachment: simplify the Attachment class

Make it a plain container around raw data and a few bits of metadata,
rather than around a whole MIME part.
---
 alot/commands/thread.py | 113 +++++++++++++++++++++++++++++-------------------
 1 file changed, 68 insertions(+), 45 deletions(-)

(limited to 'alot/commands/thread.py')

diff --git a/alot/commands/thread.py b/alot/commands/thread.py
index 0a13b868..e4334a28 100644
--- a/alot/commands/thread.py
+++ b/alot/commands/thread.py
@@ -364,11 +364,9 @@ class ForwardCommand(Command):
 
         else:  # attach original mode
             # attach original msg
-            original_mail = Message()
-            original_mail.set_type('message/rfc822')
-            original_mail['Content-Disposition'] = 'attachment'
-            original_mail.set_payload(mail.as_string(policy=email.policy.SMTP))
-            envelope.attach(Attachment(original_mail))
+            data = mail.as_string()
+            a    = Attachment(data, 'message/rfc822', None, ())
+            envelope.attach(a)
 
         # copy subject
         subject = mail.get('Subject', '')
@@ -854,48 +852,73 @@ class SaveAttachmentCommand(Command):
         self.all = all
         self.path = path
 
+    def _save_attachment(self, ui, path, attachment):
+        is_dir   = os.path.isdir(path)
+        dst_path = path
+        dst      = None
+
+        try:
+            if is_dir and not attachment.filename:
+                # generate a random filename if we don't have one
+                dst  = tempfile.NamedTemporaryFile(delete = False,
+                                                   dir = dst_path)
+                dst_path = f.name
+            else:
+                if is_dir:
+                    dst_path = os.path.join(dst_path, attachment.filename)
+
+                dst = open(dst_path, 'xb')
+
+            dst.write(attachment.data)
+            dst.close()
+        except Exception as e:
+            if dst:
+                os.remove(dst_path)
+                dst.close()
+            if isinstance(e, IOError) or isinstance(e, OSError):
+                ui.notify('Error saving attachment: %s' % str(e),
+                          priority = 'error')
+            else:
+                raise
+        else:
+            ui.notify('saved %s as: %s' % (attachment, dst_path))
+
     async def apply(self, ui):
         pcomplete = PathCompleter()
-        savedir = settings.get('attachment_prefix', '~')
+        savedir   = settings.get('attachment_prefix', '~')
+        path      = self.path
+
         if self.all:
             msg = ui.current_buffer.get_selected_message()
-            if not self.path:
-                self.path = await ui.prompt('save attachments to',
-                                            text=os.path.join(savedir, ''),
-                                            completer=pcomplete)
-            if self.path:
-                if os.path.isdir(os.path.expanduser(self.path)):
-                    for a in msg.iter_attachments():
-                        dest = a.save(self.path)
-                        name = a.get_filename()
-                        if name:
-                            ui.notify('saved %s as: %s' % (name, dest))
-                        else:
-                            ui.notify('saved attachment as: %s' % dest)
-                else:
-                    ui.notify('not a directory: %s' % self.path,
-                              priority='error')
-            else:
-                raise CommandCanceled()
+            if not path:
+                path = await ui.prompt('save attachments to',
+                                       text = os.path.join(savedir, ''),
+                                       completer = pcomplete)
+                if not path:
+                    raise CommandCanceled()
+
+            path = os.path.expanduser(path)
+            if not os.path.isdir(path):
+                ui.notify('not a directory: %s' % path,
+                          priority = 'error')
+                return
+
+            for a in msg.iter_attachments():
+                self._save_attachment(ui, path, a)
         else:  # save focussed attachment
-            attachment = ui.current_buffer.get_selected_attachment()
-            if attachment is not None:
-                filename = attachment.get_filename()
-                if not self.path:
-                    msg = 'save attachment (%s) to ' % filename
-                    initialtext = os.path.join(savedir, filename)
-                    self.path = await ui.prompt(msg,
-                                                completer=pcomplete,
-                                                text=initialtext)
-                if self.path:
-                    try:
-                        dest = attachment.save(self.path)
-                        ui.notify('saved attachment as: %s' % dest)
-                    except (IOError, OSError) as e:
-                        ui.notify(str(e), priority='error')
-                else:
+            a = ui.current_buffer.get_selected_attachment()
+            if not a:
+                return
+
+            if not path:
+                msg = 'save attachment (%s) to ' % a.filename
+                initialtext = os.path.join(savedir, '')
+                path = await ui.prompt(msg, completer = pcomplete,
+                                       text = savedir)
+                if not path:
                     raise CommandCanceled()
 
+            self._save_attachment(ui, path, a)
 
 class OpenAttachmentCommand(Command):
 
@@ -910,12 +933,12 @@ class OpenAttachmentCommand(Command):
 
     async def apply(self, ui):
         logging.info('open attachment')
-        data     = self.attachment.get_data()
-        mimetype = self.attachment.get_content_type()
-        part     = self.attachment.get_mime_representation()
-        fname    = self.attachment.get_filename()
+        data     = self.attachment.data
+        mimetype = self.attachment.content_type
+        params   = self.attachment.params
+        fname    = self.attachment.filename
 
-        h = MailcapHandler(data, mimetype, part.get_params(), fname, 'view')
+        h = MailcapHandler(data, mimetype, params, fname, 'view')
         if not h:
             ui.notify('No handler for: %s' % mimetype)
             return
-- 
cgit v1.2.3