diff --git a/afew/MailMover.py b/afew/MailMover.py index 38f25e2..62df180 100644 --- a/afew/MailMover.py +++ b/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,14 +88,14 @@ 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: @@ -128,4 +103,27 @@ def __log_move_action(self, message, source, destination, dry_run): 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) diff --git a/afew/Settings.py b/afew/Settings.py index 44d6f3c..1d4f658 100644 --- a/afew/Settings.py +++ b/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[a-z_][a-z0-9_]*)(\((?P[a-z_][a-z0-9_]*)\)|\.(?P\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['\"])(.*?)(?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 diff --git a/afew/commands.py b/afew/commands.py index 4017df6..8da57af 100644 --- a/afew/commands.py +++ b/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) diff --git a/afew/main.py b/afew/main.py index 9f91c03..6151de2 100644 --- a/afew/main.py +++ b/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.') diff --git a/afew/tests/test_mailmover.py b/afew/tests/test_mailmover.py index 9fb6986..52d01ca 100644 --- a/afew/tests/test_mailmover.py +++ b/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'))