feat: add msg attachment support (#510)

* add msg function and fix bug in eml attachment function

Author: mallorih

Diff for: CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.5.15-dev2
+
+### Features
+
+* Adds support for extracting attachments from `.msg` files
+
 ## 0.5.14-dev1
 
 ### Enhancements

Diff for: example-docs/fake-email-attachment.msg

@@ -10,7 +10,7 @@
     NarrativeText,
     Title,
 )
-from unstructured.partition.msg import partition_msg
+from unstructured.partition.msg import extract_msg_attachment_info, partition_msg
 
 DIRECTORY = pathlib.Path(__file__).parent.resolve()
 EXAMPLE_DOCS_DIRECTORY = os.path.join(DIRECTORY, "..", "..", "example-docs")
@@ -22,6 +22,15 @@
     ListItem(text="Violets are blue"),
 ]
 
+ATTACH_EXPECTED_OUTPUT = [
+    {
+        "filename": "fake-attachment.txt",
+        "extension": ".txt",
+        "file_size": "unknown",
+        "payload": b"Hey this is a fake attachment!",
+    },
+]
+
 
 def test_partition_msg_from_filename():
     filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-email.msg")
@@ -63,6 +72,13 @@ def test_partition_msg_from_file():
     assert elements == EXPECTED_MSG_OUTPUT
 
 
+def test_extract_attachment_info():
+    filename = os.path.join(DIRECTORY, "..", "..", "example-docs", "fake-email-attachment.msg")
+    attachment_info = extract_msg_attachment_info(filename)
+    assert len(attachment_info) > 0
+    assert attachment_info == ATTACH_EXPECTED_OUTPUT
+
+
 def test_partition_msg_raises_with_both_specified():
     filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-email.msg")
     with open(filename, "rb") as f, pytest.raises(ValueError):

Diff for: example-docs/fake-email-multiple-attachments.msg

@@ -1 +1 @@
-__version__ = "0.5.14-dev1"  # pragma: no cover
+__version__ = "0.5.15-dev2"  # pragma: no cover

Diff for: test_unstructured/partition/test_msg.py

@@ -123,13 +123,15 @@ def extract_attachment_info(
     output_dir: Optional[str] = None,
 ) -> List[Dict[str, str]]:
     list_attachments = []
-    attachment_info = {}
+
     for part in message.walk():
         if "content-disposition" in part:
             cdisp = part["content-disposition"].split(";")
             cdisp = [clean_extra_whitespace(item) for item in cdisp]
 
             for item in cdisp:
+                attachment_info = {}
+
                 if item.lower() == "attachment":
                     continue
                 key, value = item.split("=")

Diff for: unstructured/__version__.py

@@ -1,5 +1,5 @@
 import tempfile
-from typing import IO, List, Optional
+from typing import IO, Dict, List, Optional
 
 import msg_parser
 
@@ -67,3 +67,38 @@ def build_msg_metadata(msg_obj: msg_parser.MsOxMessage) -> ElementMetadata:
         subject=getattr(msg_obj, "subject", None),
         date=email_date,
     )
+
+
+def extract_msg_attachment_info(
+    filename: str,
+    file: Optional[IO] = None,
+    output_dir: Optional[str] = None,
+) -> List[Dict[str, str]]:
+    exactly_one(filename=filename, file=file)
+
+    if filename is not None:
+        msg_obj = msg_parser.MsOxMessage(filename)
+    elif file is not None:
+        tmp = tempfile.NamedTemporaryFile(delete=False)
+        tmp.write(file.read())
+        tmp.close()
+        msg_obj = msg_parser.MsOxMessage(tmp.name)
+
+    list_attachments = []
+
+    for attachment in msg_obj.attachments:
+        attachment_info = {}
+
+        attachment_info["filename"] = attachment.AttachLongFilename
+        attachment_info["extension"] = attachment.AttachExtension
+        attachment_info["file_size"] = attachment.AttachmentSize
+        attachment_info["payload"] = attachment.data
+
+        list_attachments.append(attachment_info)
+
+        if output_dir is not None:
+            filename = output_dir + "/" + attachment_info["filename"]
+            with open(filename, "wb") as f:
+                f.write(attachment.data)
+
+    return list_attachments