Merge pull request #13 from loadsmart/recursivePartLookup

Support recursive parts for attachment lookup
loadsmart · Jan 30, 2020 · 036d991 · 036d991
2 parents 41a7489 + 6a723d1
commit 036d991
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 11 deletions.
diff --git a/gmail_wrapper/entities.py b/gmail_wrapper/entities.py
@@ -106,14 +106,24 @@ def date(self):
         date_in_seconds = int(self._raw.get("internalDate")) / ms_in_seconds
         return datetime.utcfromtimestamp(date_in_seconds)
 
+    def _extract_attachments_from_parts(self, parts):
+        attachments = []
+        for part in parts:
+            if part.get("filename") and part.get("body", {}).get("attachmentId"):
+                attachments.append(Attachment(self.id, self._client, part))
+            if part.get("parts"):
+                attachments += self._extract_attachments_from_parts(part.get("parts"))
+
+        return attachments
+
     @property
     def attachments(self):
         parts = self._payload.get("parts")
-        return [
-            Attachment(self.id, self._client, part)
-            for part in (parts if parts else [])
-            if part["filename"] and part["body"] and part["body"].get("attachmentId")
-        ]
+
+        if not parts:
+            return []
+
+        return self._extract_attachments_from_parts(parts)
 
     def modify(self, add_labels=None, remove_labels=None):
         self._raw = self._client.modify_raw_message(

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -52,9 +52,9 @@ def raw_complete_message():
             "parts": [
                 {
                     "body": {
-                        "data": "I am caught up in a meeting. Can you run an urgent errand for me?",
+                        "data": "I am Dr. Bakare Tunde, the cousin of Nigerian Astronaut, Air Force Major Abacha Tunde. He was the first African in space when he made a secret flight to the Salyut 6 space station in 1979.",
                         "attachmentId": "CCX456",
-                        "size": 65,
+                        "size": 188,
                     },
                     "mimeType": "text/plain",
                     "partId": "BB789",
@@ -66,10 +66,29 @@ def raw_complete_message():
                     "partId": "BB790",
                     "filename": "fox.txt",
                 },
+                {
+                    "body": {"size": 0},
+                    "mimeType": "multipart/mixed",
+                    "partId": "BB791",
+                    "parts": [
+                        {
+                            "body": {"data": "someRandomDataHere", "size": 18},
+                            "mimeType": "text/html",
+                            "partId": "BB791.0",
+                            "filename": "",
+                        },
+                        {
+                            "body": {"attachmentId": "CCX458", "size": 95},
+                            "mimeType": "application/pdf",
+                            "partId": "BB791.1",
+                            "filename": "tigers.pdf",
+                        }
+                    ],
+                },
             ],
         },
-        "snippet": "I am caught up in a meeting. Can you run an urgent errand for me?",
-        "sizeEstimate": 125,
+        "snippet": "I am Dr. Bakare Tunde, the cousin of Nigerian Astronaut, Air Force Major Abacha Tunde.",
+        "sizeEstimate": 361,
         "threadId": "AA121212",
         "labelIds": ["phishing"],
         "id": "123AAB",

diff --git a/tests/test_entities.py b/tests/test_entities.py
@@ -77,9 +77,11 @@ def test_it_does_not_generate_attachment_objects_for_attachments_without_id(
                 "mimeType": "application/pdf",
             }
         )
-        assert len(complete_message.attachments) == 1
+        assert len(complete_message.attachments) == 2
         assert complete_message.attachments[0].id
-        assert complete_message.attachments[0].filename != "invalid.pdf"
+        assert complete_message.attachments[0].filename == "fox.txt"
+        assert complete_message.attachments[1].id
+        assert complete_message.attachments[1].filename == "tigers.pdf"
 
     def test_it_returns_empty_list_if_no_attachments(
         self, client, raw_complete_message