Abstract out MailMover from mail selection

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/Settings.py

@@ -23,16 +23,14 @@
 
 # All the values for keys listed here are interpreted as ;-delimited lists
 value_is_a_list = ['tags', 'tags_blacklist']
-mail_mover_section = 'MailMover'
+folder_mail_mover_section = 'MailMover'
 
 section_re = re.compile(r'^(?P<name>[a-z_][a-z0-9_]*)(\((?P<parent_class>[a-z_][a-z0-9_]*)\)|\.(?P<index>\d+))?$', re.I)
-
-
 def get_filter_chain(database):
     filter_chain = []
 
     for section in settings.sections():
-        if section == 'global' or section == mail_mover_section:
+        if section in ['global', folder_mail_mover_section]:
             continue
 
         match = section_re.match(section)
@@ -64,17 +62,29 @@ def get_filter_chain(database):
 
     return filter_chain
 
+def get_mail_move_kind():
+    return settings.get('global', 'mail_mover_kind', fallback='folder')
+
+def get_mail_move_section(kind):
+    if kind == 'folder':
+        return folder_mail_mover_section
+
+def get_mail_move_rules(kind):
+    section = get_mail_move_section(kind)
+    if kind == 'query':
+        rule_id_key = 'rules'
+    else:
+        rule_id_key = 'folders'
 
-def get_mail_move_rules():
     rule_pattern = re.compile(r"'(.+?)':((?P<quote>['\"])(.*?)(?P=quote)|\S+)")
-    if settings.has_option(mail_mover_section, 'folders'):
+    if settings.has_option(section, rule_id_key):
         all_rules = collections.OrderedDict()
 
-        for folder in shlex.split(settings.get(mail_mover_section, 'folders')):
-            if settings.has_option(mail_mover_section, folder):
+        for folder in shlex.split(settings.get(section, rule_id_key)):
+            if settings.has_option(section, folder):
                 rules = collections.OrderedDict()
                 raw_rules = re.findall(rule_pattern,
-                                       settings.get(mail_mover_section, folder))
+                                       settings.get(section, folder))
                 for rule in raw_rules:
                     query = rule[0]
                     destination = rule[3] or rule[1]
@@ -87,16 +97,14 @@ def get_mail_move_rules():
     else:
         raise NameError("No folders defined to move mails from.")
 
-
-def get_mail_move_age():
+def get_mail_move_age(section):
     max_age = 0
-    if settings.has_option(mail_mover_section, 'max_age'):
-        max_age = settings.get(mail_mover_section, 'max_age')
+    if settings.has_option(section, 'max_age'):
+        max_age = settings.get(section, 'max_age')
     return max_age
 
-
-def get_mail_move_rename():
+def get_mail_move_rename(section):
     rename = False
-    if settings.has_option(mail_mover_section, 'rename'):
-        rename = settings.get(mail_mover_section, 'rename').lower() == 'true'
+    if settings.has_option(section, 'rename'):
+        rename = settings.get(section, 'rename').lower() == 'true'
     return rename

afew/commands.py

@@ -10,7 +10,8 @@
 from afew.main import main as inner_main
 from afew.FilterRegistry import all_filters
 from afew.Settings import user_config_dir, get_filter_chain, \
-    get_mail_move_rules, get_mail_move_age, get_mail_move_rename
+    get_mail_move_kind, get_mail_move_section, get_mail_move_rules, \
+    get_mail_move_age, get_mail_move_rename
 from afew.NotmuchSettings import read_notmuch_settings, get_notmuch_new_query
 from afew.version import version
 
@@ -136,9 +137,12 @@ def main():
         __import__(file_name[:-3], level=0)
 
     if args.move_mails:
-        args.mail_move_rules = get_mail_move_rules()
-        args.mail_move_age = get_mail_move_age()
-        args.mail_move_rename = get_mail_move_rename()
+        args.mail_move_kind = get_mail_move_kind()
+        section = get_mail_move_section(args.mail_move_kind)
+
+        args.mail_move_rules = get_mail_move_rules(args.mail_move_kind)
+        args.mail_move_age = get_mail_move_age(section)
+        args.mail_move_rename = get_mail_move_rename(section)
 
     with Database() as database:
         configured_filter_chain = get_filter_chain(database)

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
@@ -24,9 +24,18 @@ def main(options, database, query_string):
         watch_for_new_files(options, database,
                             quick_find_dirs_hack(database.db_path))
     elif options.move_mails:
+        if options.mail_move_kind == 'folder':
+            mover_class = FolderMailMover
+        else:
+            sys.exit('Mail mover kind {:r} is not recognized'.format(options.mail_move_kind))
+
         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 = mover_class(
+                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

@@ -42,8 +42,7 @@ def create_mail(msg, maildir, notmuch_db, tags, old=False):
     return (stripped_msgid, msg)
 
 
-@freeze_time("2019-01-30 12:00:00")
-class TestMailMover(unittest.TestCase):
+class MailMoverTestBaseClass:
     def setUp(self):
         self.test_dir = tempfile.mkdtemp()
 
@@ -62,11 +61,30 @@ 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.
+
+    def tearDown(self):
+        shutil.rmtree(self.test_dir)
+
+
+    @staticmethod
+    def get_folder_content(db, folder):
+        return {
+            (os.path.basename(msg.get_message_id()), msg.get_part(1).decode())
+            for msg in db.do_query('folder:{}'.format(folder)).search_messages()
+        }
+
+
+
+@freeze_time("2019-01-30 12:00:00")
+class TestFolderMailMover(MailMoverTestBaseClass, unittest.TestCase):
+    def __init__(self, *args, **kwargs):
+        super(TestFolderMailMover, self).__init__(*args, **kwargs)
+
+        # 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': {
@@ -84,18 +102,6 @@ def setUp(self):
         }
 
 
-    def tearDown(self):
-        shutil.rmtree(self.test_dir)
-
-
-    @staticmethod
-    def get_folder_content(db, folder):
-        return {
-            (os.path.basename(msg.get_message_id()), msg.get_part(1).decode())
-            for msg in db.do_query('folder:{}'.format(folder)).search_messages()
-        }
-
-
     def test_all_rule_cases(self):
         from afew import MailMover
 
@@ -121,11 +127,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 +152,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'))