diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..988937c --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.pythonPath": "/usr/local/bin/python3" +} \ No newline at end of file diff --git a/imbox/messages.py b/imbox/messages.py index bc5b3c0..d11bec8 100644 --- a/imbox/messages.py +++ b/imbox/messages.py @@ -1,6 +1,7 @@ import datetime import logging +from imbox.query import build_search_query from imbox.parser import fetch_email_by_uid @@ -11,6 +12,7 @@ class Messages: IMAP_ATTRIBUTE_LOOKUP = { 'unread': '(UNSEEN)', + 'flagged': '(FLAGGED)', 'unflagged': '(UNFLAGGED)', 'sent_from': '(FROM "{}")', 'sent_to': '(TO "{}")', @@ -18,7 +20,7 @@ class Messages: 'date__lt': '(BEFORE "{}")', 'date__on': '(ON "{}")', 'subject': '(SUBJECT "{}")', - 'uid_range': '(UID {})', + 'uid__range': '(UID {})', } FOLDER_LOOKUP = {} @@ -41,25 +43,12 @@ class Messages: parser_policy=self.parser_policy) def _query_uids(self, **kwargs): - query_ = self._build_search_query(**kwargs) + query_ = build_search_query(self.IMAP_ATTRIBUTE_LOOKUP, **kwargs) _, data = self.connection.uid('search', None, query_) if data[0] is None: return [] return data[0].split() - def _build_search_query(self, **kwargs): - query = [] - for name, value in kwargs.items(): - if value is not None: - if isinstance(value, datetime.date): - value = value.strftime('%d-%b-%Y') - query.append(self.IMAP_ATTRIBUTE_LOOKUP[name].format(value)) - - if query: - return " ".join(query) - - return "(ALL)" - def _fetch_email_list(self): for uid in self._uid_list: yield uid, self._fetch_email(uid) diff --git a/imbox/query.py b/imbox/query.py new file mode 100644 index 0000000..8ad3bb5 --- /dev/null +++ b/imbox/query.py @@ -0,0 +1,17 @@ +import datetime + + +def build_search_query(imap_attribute_lookup, **kwargs): + query = [] + for name, value in kwargs.items(): + if value is not None: + if isinstance(value, datetime.date): + value = value.strftime('%d-%b-%Y') + if isinstance(value, str) and '"' in value: + value = value.replace('"', "'") + query.append(imap_attribute_lookup[name].format(value)) + + if query: + return " ".join(query) + + return "(ALL)" diff --git a/imbox/vendors/gmail.py b/imbox/vendors/gmail.py index f6cf3c8..85d05fa 100644 --- a/imbox/vendors/gmail.py +++ b/imbox/vendors/gmail.py @@ -1,4 +1,5 @@ from imbox.messages import Messages +from imbox.vendors.helpers import merge_two_dicts class GmailMessages(Messages): @@ -21,15 +22,18 @@ class GmailMessages(Messages): 'trash': '"[Gmail]/Trash"', } + GMAIL_IMAP_ATTRIBUTE_LOOKUP_DIFF = { + 'subject': '(X-GM-RAW "subject:\'{}\'")', + 'label': '(X-GM-LABELS "{}")', + 'raw': '(X-GM-RAW "{}")' + } + def __init__(self, connection, parser_policy, **kwargs): - self.IMAP_ATTRIBUTE_LOOKUP = {**self.IMAP_ATTRIBUTE_LOOKUP, **{'subject': '(X-GM-RAW "{}")', - 'label': '(X-GM-LABELS "{}")', - 'raw': '(X-GM-RAW "{}")', - } - } + self.IMAP_ATTRIBUTE_LOOKUP = merge_two_dicts(self.IMAP_ATTRIBUTE_LOOKUP, + self.GMAIL_IMAP_ATTRIBUTE_LOOKUP_DIFF) super().__init__(connection, parser_policy, **kwargs) diff --git a/imbox/vendors/helpers.py b/imbox/vendors/helpers.py new file mode 100644 index 0000000..8b3f0a6 --- /dev/null +++ b/imbox/vendors/helpers.py @@ -0,0 +1,6 @@ + +def merge_two_dicts(x, y): + """from https://stackoverflow.com/a/26853961/4386191""" + z = x.copy() # start with x's keys and values + z.update(y) # modifies z with y's keys and values & returns None + return z diff --git a/tests/query_tests.py b/tests/query_tests.py index 1bf7ba7..af9158d 100644 --- a/tests/query_tests.py +++ b/tests/query_tests.py @@ -1,54 +1,83 @@ -import unittest -from imbox.query import build_search_query from datetime import date +import unittest + +from imbox.query import build_search_query +from imbox.messages import Messages +from imbox.vendors.helpers import merge_two_dicts +from imbox.vendors.gmail import GmailMessages + +IMAP_ATTRIBUTE_LOOKUP = Messages.IMAP_ATTRIBUTE_LOOKUP +GMAIL_ATTRIBUTE_LOOKUP = merge_two_dicts(IMAP_ATTRIBUTE_LOOKUP, + GmailMessages.GMAIL_IMAP_ATTRIBUTE_LOOKUP_DIFF) class TestQuery(unittest.TestCase): def test_all(self): - res = build_search_query() + res = build_search_query(IMAP_ATTRIBUTE_LOOKUP) self.assertEqual(res, "(ALL)") + def test_subject(self): + + res = build_search_query(IMAP_ATTRIBUTE_LOOKUP, subject='hi') + self.assertEqual(res, '(SUBJECT "hi")') + + res = build_search_query(GMAIL_ATTRIBUTE_LOOKUP, subject='hi') + self.assertEqual(res, '(X-GM-RAW "subject:\'hi\'")') + def test_unread(self): - res = build_search_query(unread=True) + res = build_search_query(IMAP_ATTRIBUTE_LOOKUP, unread=True) self.assertEqual(res, "(UNSEEN)") def test_unflagged(self): - res = build_search_query(unflagged=True) + res = build_search_query(IMAP_ATTRIBUTE_LOOKUP, unflagged=True) self.assertEqual(res, "(UNFLAGGED)") def test_flagged(self): - res = build_search_query(flagged=True) + res = build_search_query(IMAP_ATTRIBUTE_LOOKUP, flagged=True) self.assertEqual(res, "(FLAGGED)") def test_sent_from(self): - res = build_search_query(sent_from='test@example.com') + res = build_search_query( + IMAP_ATTRIBUTE_LOOKUP, sent_from='test@example.com') self.assertEqual(res, '(FROM "test@example.com")') def test_sent_to(self): - res = build_search_query(sent_to='test@example.com') + res = build_search_query( + IMAP_ATTRIBUTE_LOOKUP, sent_to='test@example.com') self.assertEqual(res, '(TO "test@example.com")') def test_date__gt(self): - res = build_search_query(date__gt=date(2014, 12, 31)) + res = build_search_query( + IMAP_ATTRIBUTE_LOOKUP, date__gt=date(2014, 12, 31)) self.assertEqual(res, '(SINCE "31-Dec-2014")') def test_date__lt(self): - res = build_search_query(date__lt=date(2014, 1, 1)) + res = build_search_query( + IMAP_ATTRIBUTE_LOOKUP, date__lt=date(2014, 1, 1)) self.assertEqual(res, '(BEFORE "01-Jan-2014")') def test_date__on(self): - res = build_search_query(date__on=date(2014, 1, 1)) + res = build_search_query( + IMAP_ATTRIBUTE_LOOKUP, date__on=date(2014, 1, 1)) self.assertEqual(res, '(ON "01-Jan-2014")') def test_uid__range(self): - res = build_search_query(uid__range='1000:*') + res = build_search_query(IMAP_ATTRIBUTE_LOOKUP, uid__range='1000:*') self.assertEqual(res, '(UID 1000:*)') + + def test_gmail_raw(self): + res = build_search_query(GMAIL_ATTRIBUTE_LOOKUP, raw='has:attachment subject:"hey"') + self.assertEqual(res, '(X-GM-RAW "has:attachment subject:\'hey\'")') + + def test_gmail_label(self): + res = build_search_query(GMAIL_ATTRIBUTE_LOOKUP, label='finance') + self.assertEqual(res, '(X-GM-LABELS "finance")')