MailMover: Split finding mails to separate class

This allows implementing multiple move strategies.

Signed-off-by: Johannes Löthberg <[email protected]>

afew/MailMover.py

@@ -5,6 +5,7 @@
 import os
 import shutil
 import uuid
+from abc import ABC, abstractmethod
 from datetime import date, datetime, timedelta
 from subprocess import check_call, CalledProcessError, DEVNULL
 
@@ -14,21 +15,9 @@
 from afew.utils import get_message_summary
 
 
-class MailMover(Database):
-    """
-    Move mail files matching a given notmuch query into a target maildir folder.
-    """
-
-    def __init__(self, max_age=0, rename=False, dry_run=False, notmuch_args='', quiet=False):
-        super().__init__()
-        self.db = notmuch.Database(self.db_path)
-        self.query = 'folder:"{folder}" AND {subquery}'
-        if max_age:
-            days = timedelta(int(max_age))
-            start = date.today() - days
-            now = datetime.now()
-            self.query += ' AND {start}..{now}'.format(start=start.strftime('%s'),
-                                                       now=now.strftime('%s'))
+class AbstractMailMover(ABC):
+    def __init__(self, rename=False, dry_run=False, notmuch_args='', quiet=False):
+        self.db = Database()
         self.dry_run = dry_run
         self.rename = rename
         self.notmuch_args = notmuch_args
@@ -48,61 +37,47 @@ def get_new_name(self, fname, destination):
             basename = str(uuid.uuid1()) + flagpart
         return os.path.join(destination, submaildir, basename)
 
-    def move(self, maildir, rules):
-        """
-        Move mails in folder maildir according to the given rules.
-        """
-        # identify and move messages
-        logging.info("checking mails in '{}'".format(maildir))
-        to_delete_fnames = []
+
+    def move(self, rule_name, rules):
+        logging.info("Processing rule '{}'".format(rule_name))
+
         moved = False
-        for query in rules.keys():
-            destination = '{}/{}/'.format(self.db_path, rules[query])
-            main_query = self.query.format(
-                folder=maildir.replace("\"", "\\\""), subquery=query)
-            logging.debug("query: {}".format(main_query))
-            messages = notmuch.Query(self.db, main_query).search_messages()
-            for message in messages:
-                # a single message (identified by Message-ID) can be in several
-                # places; only touch the one(s) that exists in this maildir
-                all_message_fnames = message.get_filenames()
-                to_move_fnames = [name for name in all_message_fnames
-                                  if maildir in name]
-                if not to_move_fnames:
-                    continue
+        fnames_to_delete = []
+        for query, dest_maildir in rules.items():
+            destination = '{}/{}/'.format(self.db.db_path, dest_maildir)
+            for (message, fname) in self.find_matching(rule_name, query):
                 moved = True
-                self.__log_move_action(message, maildir, rules[query],
-                                       self.dry_run)
-                for fname in to_move_fnames:
-                    if self.dry_run:
-                        continue
-                    try:
-                        shutil.copy2(fname, self.get_new_name(fname, destination))
-                        to_delete_fnames.append(fname)
-                    except shutil.SameFileError:
-                        logging.warn("trying to move '{}' onto itself".format(fname))
+                self.__log_move_action(message, dest_maildir)
+                try:
+                    shutil.copy2(fname, self.get_new_name(fname, destination))
+                    fnames_to_delete.append(fname)
+                except shutil.SameFileError:
+                    logging.warn("trying to move '{}' onto itself".format(fname))
+                    continue
+                except shutil.Error as e:
+                    # this is ugly, but shutil does not provide more
+                    # finely individuated errors
+                    if str(e).endswith("already exists"):
                         continue
-                    except shutil.Error as e:
-                        # this is ugly, but shutil does not provide more
-                        # finely individuated errors
-                        if str(e).endswith("already exists"):
-                            continue
-                        else:
-                            raise
+                    else:
+                        raise
+
+        # close database after we're done using it
+        self.db.close()
 
         # remove mail from source locations only after all copies are finished
-        for fname in set(to_delete_fnames):
+        for fname in set(fnames_to_delete):
             os.remove(fname)
 
         # update notmuch database
         if not self.dry_run:
             if moved:
                 logging.info("updating database")
-                self.__update_db(maildir)
+                self.__update_db()
         else:
             logging.info("Would update database")
 
-    def __update_db(self, maildir):
+    def __update_db(self):
         """
         Update the database after mail files have been moved in the filesystem.
         """
@@ -113,19 +88,42 @@ def __update_db(self, maildir):
                 check_call(['notmuch', 'new'] + self.notmuch_args.split())
         except CalledProcessError as err:
             logging.error("Could not update notmuch database "
-                          "after syncing maildir '{}': {}".format(maildir, err))
+                          "after syncing: {}".format(err))
             raise SystemExit
 
-    def __log_move_action(self, message, source, destination, dry_run):
-        '''
+    def __log_move_action(self, message, destination):
+        """
         Report which mails have been identified for moving.
-        '''
-        if not dry_run:
+        """
+        if not self.dry_run:
             level = logging.DEBUG
             prefix = 'moving mail'
         else:
             level = logging.INFO
             prefix = 'I would move mail'
         logging.log(level, prefix)
         logging.log(level, "    {}".format(get_message_summary(message).encode('utf8')))
-        logging.log(level, "from '{}' to '{}'".format(source, destination))
+        logging.log(level, "to '{}'".format(destination))
+
+
+class FolderMailMover(AbstractMailMover):
+    def __init__(self, max_age=0, *args, **kwargs):
+        super(FolderMailMover, self).__init__(*args, **kwargs)
+        self.query = 'folder:"{folder}" AND {subquery}'
+        if max_age:
+            days = timedelta(int(max_age))
+            start = date.today() - days
+            now = datetime.now()
+            self.query += ' AND {start}..{now}'.format(start=start.strftime('%s'),
+                                                       now=now.strftime('%s'))
+
+    def find_matching(self, maildir, query):
+        main_query = self.query.format(
+            folder=maildir.replace('"', '\\"'),
+            subquery=query
+        )
+        for message in self.db.do_query(main_query).search_messages():
+            # a single message (identified by Message-ID) can be in several
+            # places; only touch the one(s) that exists in this maildir
+            for fname in [fname for fname in message.get_filenames() if maildir in fname]:
+                yield (message, fname)

afew/main.py

@@ -3,7 +3,7 @@
 
 import sys
 
-from afew.MailMover import MailMover
+from afew.MailMover import FolderMailMover
 
 try:
     from .files import watch_for_new_files, quick_find_dirs_hack
@@ -25,8 +25,12 @@ def main(options, database, query_string):
                             quick_find_dirs_hack(database.db_path))
     elif options.move_mails:
         for maildir, rules in options.mail_move_rules.items():
-            mover = MailMover(options.mail_move_age, options.mail_move_rename, options.dry_run, options.notmuch_args)
+            mover = FolderMailMover(
+                max_age=options.mail_move_age,
+                rename=options.mail_move_rename,
+                dry_run=options.dry_run,
+                notmuch_args=options.notmuch_args
+            )
             mover.move(maildir, rules)
-            mover.close()
     else:
         sys.exit('Weird... please file a bug containing your command line.')

afew/tests/test_mailmover.py

@@ -43,7 +43,7 @@ def create_mail(msg, maildir, notmuch_db, tags, old=False):
 
 
 @freeze_time("2019-01-30 12:00:00")
-class TestMailMover(unittest.TestCase):
+class TestFolderMailMover(unittest.TestCase):
     def setUp(self):
         self.test_dir = tempfile.mkdtemp()
 
@@ -62,11 +62,11 @@ def setUp(self):
         self.archive = self.root.add_folder('archive')
         self.spam = self.root.add_folder('spam')
 
-        # Dict of rules that are passed to MailMover.
+        # Dict of rules that are passed to FolderMailMover.
         #
         # The top level key represents a particular mail directory to work on.
         #
-        # The second level key is the notmuch query that MailMover will execute,
+        # The second level key is the notmuch query that FolderMailMover will execute,
         # and its value is the directory to move the matching emails to.
         self.rules = {
             '.inbox': {
@@ -121,11 +121,10 @@ def test_all_rule_cases(self):
                 create_mail('In spam, tagged archive, spam\n', self.spam, db, ['archive', 'spam']),
             ])
 
-        mover = MailMover.MailMover(quiet=True)
+        mover = MailMover.FolderMailMover(quiet=True)
         mover.move('.inbox', self.rules['.inbox'])
         mover.move('.archive', self.rules['.archive'])
         mover.move('.spam', self.rules['.spam'])
-        mover.close()
 
         with Database() as db:
             self.assertEqual(expect_inbox, self.get_folder_content(db, '.inbox'))
@@ -147,11 +146,10 @@ def test_max_age(self):
 
             expect_spam = set([])
 
-        mover = MailMover.MailMover(max_age=15, quiet=True)
+        mover = MailMover.FolderMailMover(max_age=15, quiet=True)
         mover.move('.inbox', self.rules['.inbox'])
         mover.move('.archive', self.rules['.archive'])
         mover.move('.spam', self.rules['.spam'])
-        mover.close()
 
         with Database() as db:
             self.assertEqual(expect_inbox, self.get_folder_content(db, '.inbox'))