Added ability to attach files
authorDmitriy Morozov <morozov@cs.duke.edu>
Mon, 28 Apr 2008 18:07:55 -0400
changeset 26 4574d2d34009
parent 25 a13239888b62
child 27 6ab60ee8b151
Added ability to attach files
artemis.py
--- a/artemis.py	Tue Apr 22 07:23:05 2008 -0400
+++ b/artemis.py	Mon Apr 28 18:07:55 2008 -0400
@@ -5,6 +5,14 @@
 from mercurial import hg, util
 from mercurial.i18n import _
 import os, time, random, mailbox, glob, socket, ConfigParser
+import mimetypes
+from email import encoders
+from email.generator import Generator
+from email.mime.audio import MIMEAudio
+from email.mime.base import MIMEBase
+from email.mime.image import MIMEImage
+from email.mime.multipart import MIMEMultipart
+from email.mime.text import MIMEText
 
 
 state = {'new': 'new', 'fixed': 'fixed'}
@@ -31,11 +39,7 @@
 
     issues = glob.glob(os.path.join(issues_path, '*'))
 
-    # Create missing dirs
-    for i in issues:
-        for d in maildir_dirs:
-            path = os.path.join(issues_path,i,d)
-            if not os.path.exists(path): os.mkdir(path)
+    _create_all_missing_dirs(issues_path, issues)
 
     # Process filter
     if opts['filter']:
@@ -65,7 +69,7 @@
                                           mbox[root]['Subject']))
 
 
-def iadd(ui, repo, id = None, comment = 0):
+def iadd(ui, repo, id = None, comment = 0, **opts):
     """Adds a new issue, or comment to an existing issue ID or its comment COMMENT"""
 
     comment = int(comment)
@@ -79,9 +83,7 @@
         if not issue_fn:
             ui.warn('No such issue\n')
             return
-        for d in maildir_dirs:
-            path = os.path.join(issues_path,issue_id,d)
-            if not os.path.exists(path): os.mkdir(path)
+        _create_missing_dirs(issues_path, issue_id)
 
     user = ui.username()
 
@@ -101,7 +103,10 @@
 
     # Create the message
     msg = mailbox.MaildirMessage(issue)
-    #msg.set_from('artemis', True)
+    if opts['attach']:
+        outer = _attach_files(msg, opts['attach'])
+    else:
+        outer = msg
 
     # Pick random filename
     if not id:
@@ -119,13 +124,13 @@
         ui.warn('No such comment number in mailbox, commenting on the issue itself\n')
 
     if not id:
-        msg.add_header('Message-Id', "<%s-0-artemis@%s>" % (issue_id, socket.gethostname()))
+        outer.add_header('Message-Id', "<%s-0-artemis@%s>" % (issue_id, socket.gethostname()))
     else:
         root = keys[0]
-        msg.add_header('Message-Id', "<%s-%s-artemis@%s>" % (issue_id, _random_id(), socket.gethostname()))
-        msg.add_header('References', mbox[(comment < len(mbox) and keys[comment]) or root]['Message-Id'])
-        msg.add_header('In-Reply-To', mbox[(comment < len(mbox) and keys[comment]) or root]['Message-Id'])
-    repo.add([issue_fn[(len(repo.root)+1):] + '/new/'  + mbox.add(msg)])   # +1 for the trailing /
+        outer.add_header('Message-Id', "<%s-%s-artemis@%s>" % (issue_id, _random_id(), socket.gethostname()))
+        outer.add_header('References', mbox[(comment < len(mbox) and keys[comment]) or root]['Message-Id'])
+        outer.add_header('In-Reply-To', mbox[(comment < len(mbox) and keys[comment]) or root]['Message-Id'])
+    repo.add([issue_fn[(len(repo.root)+1):] + '/new/'  + mbox.add(outer)])   # +1 for the trailing /
     mbox.close()
 
     # If adding issue, add the new mailbox to the repository
@@ -140,10 +145,7 @@
     issue, id = _find_issue(ui, repo, id)
     if not issue: return
     
-    # Create missing dirs
-    for d in maildir_dirs:
-        path = os.path.join(repo.root,issues_dir,issue,d)
-        if not os.path.exists(path): os.mkdir(path)
+    _create_missing_dirs(os.path.join(repo.root, issues_dir), issue)
 
     mbox = mailbox.Maildir(issue, factory=mailbox.MaildirMessage)
 
@@ -159,6 +161,25 @@
 
     _show_mbox(ui, mbox, comment)
 
+    if opts['extract']:
+        attachment_numbers = map(int, opts['extract'])
+        keys = _order_keys_date(mbox)
+        msg = mbox[keys[comment]]
+        counter = 1
+        for part in msg.walk():
+            ctype = part.get_content_type()
+            maintype, subtype = ctype.split('/', 1)
+            if maintype == 'multipart' or ctype == 'text/plain': continue
+            if counter in attachment_numbers:
+                filename = part.get_filename()
+                if not filename:
+                    ext = mimetypes.guess_extension(part.get_content_type()) or ''
+                    filename = 'attachment-%03d%s' % (counter, ext)
+                fp = open(filename, 'wb')
+                fp.write(part.get_payload(decode = True))
+                fp.close()
+            counter += 1
+
 
 def iupdate(ui, repo, id, **opts):
     """Update properties of issue ID"""
@@ -166,11 +187,7 @@
     issue, id = _find_issue(ui, repo, id)
     if not issue: return
     
-    # Create missing dirs
-    for d in maildir_dirs:
-        path = os.path.join(repo.root,issues_dir,issue,d)
-        if not os.path.exists(path): os.mkdir(path)
-
+    _create_missing_dirs(os.path.join(repo.root, issues_dir), issue_id)
 
     properties = _get_properties(opts['property'])
 
@@ -236,7 +253,17 @@
         if 'Date' in message: ui.write('Date: %s\n' % message['Date'])
         if 'Subject' in message: ui.write('Subject: %s\n' % message['Subject'])
         if 'State' in message: ui.write('State: %s\n' % message['State'])
-        ui.write('\n' + message.get_payload().strip() + '\n')
+        counter = 1
+        for part in message.walk():
+            ctype = part.get_content_type()
+            maintype, subtype = ctype.split('/', 1)
+            if maintype == 'multipart': continue
+            if ctype == 'text/plain':
+                ui.write('\n' + part.get_payload().strip() + '\n')
+            else:
+                filename = part.get_filename()
+                ui.write('\n' + '%d: Attachment [%s, %s]: %s' % (counter, ctype, _humanreadable(len(part.get_payload())), filename) + '\n')
+                counter += 1
 
 def _show_mbox(ui, mbox, comment):
     # Output the issue (or comment)
@@ -297,6 +324,59 @@
 def _random_id():
     return "%x" % random.randint(2**63, 2**64-1)
 
+def _create_missing_dirs(issues_path, issue):
+    for d in maildir_dirs:
+        path = os.path.join(issues_path,issue,d)
+        if not os.path.exists(path): os.mkdir(path)
+
+def _create_all_missing_dirs(issues_path, issues):
+    for i in issues:
+        _create_missing_dirs(issues_path, i)
+
+def _humanreadable(size):
+    if size > 1024*1024:
+        return '%5.1fM' % (float(size) / (1024*1024))
+    elif size > 1024:
+        return '%5.1fK' % (float(size) / 1024)
+    else:
+        return '%dB' % size
+
+def _attach_files(msg, filenames):
+    outer = MIMEMultipart()
+    for k in msg.keys(): outer[k] = msg[k]
+    outer.attach(MIMEText(msg.get_payload()))
+
+    for filename in filenames:
+        ctype, encoding = mimetypes.guess_type(filename)
+        if ctype is None or encoding is not None:
+            # No guess could be made, or the file is encoded (compressed), so
+            # use a generic bag-of-bits type.
+            ctype = 'application/octet-stream'
+        maintype, subtype = ctype.split('/', 1)
+        if maintype == 'text':
+            fp = open(filename)
+            # Note: we should handle calculating the charset
+            attachment = MIMEText(fp.read(), _subtype=subtype)
+            fp.close()
+        elif maintype == 'image':
+            fp = open(filename, 'rb')
+            attachment = MIMEImage(fp.read(), _subtype=subtype)
+            fp.close()
+        elif maintype == 'audio':
+            fp = open(filename, 'rb')
+            attachment = MIMEAudio(fp.read(), _subtype=subtype)
+            fp.close()
+        else:
+            fp = open(filename, 'rb')
+            attachment = MIMEBase(maintype, subtype)
+            attachment.set_payload(fp.read())
+            fp.close()
+            # Encode the payload using Base64
+            encoders.encode_base64(attachment)
+        # Set the filename parameter
+        attachment.add_header('Content-Disposition', 'attachment', filename=filename)
+        outer.attach(attachment)
+    return outer
 
 cmdtable = {
     'ilist':    (ilist,
@@ -308,10 +388,12 @@
                   ('f', 'filter', '', 'restrict to pre-defined filter (in %s/%s*)' % (issues_dir, filter_prefix))],
                  _('hg ilist [OPTIONS]')),
     'iadd':       (iadd, 
-                 [],
+                 [('a', 'attach', [],
+                   'attach file(s) (e.g., -a filename1 -a filename2)')], 
                  _('hg iadd [ID] [COMMENT]')),
     'ishow':      (ishow,
-                 [('a', 'all', None, 'list all comments')],
+                 [('a', 'all', None, 'list all comments'),
+                  ('x', 'extract', [], 'extract attachments')],
                  _('hg ishow [OPTIONS] ID [COMMENT]')),
     'iupdate':    (iupdate,
                  [('p', 'property', [],