Merge pull request #38 from VillePuuska/support-comments-in-docx

Add passing style_map kwarg to Mammoth when converting docx to allow keeping comments

src/markitdown/_markitdown.py

@@ -492,7 +492,9 @@ def convert(self, local_path, **kwargs) -> Union[None, DocumentConverterResult]:
 
         result = None
         with open(local_path, "rb") as docx_file:
-            result = mammoth.convert_to_html(docx_file)
+            style_map = kwargs.get("style_map", None)
+
+            result = mammoth.convert_to_html(docx_file, style_map=style_map)
             html_content = result.value
             result = self._convert(html_content)
 
@@ -999,6 +1001,7 @@ def __init__(
         requests_session: Optional[requests.Session] = None,
         mlm_client: Optional[Any] = None,
         mlm_model: Optional[Any] = None,
+        style_map: Optional[str] = None,
     ):
         if requests_session is None:
             self._requests_session = requests.Session()
@@ -1007,6 +1010,7 @@ def __init__(
 
         self._mlm_client = mlm_client
         self._mlm_model = mlm_model
+        self._style_map = style_map
 
         self._page_converters: List[DocumentConverter] = []
 
@@ -1184,6 +1188,9 @@ def _convert(
                 # Add the list of converters for nested processing
                 _kwargs["_parent_converters"] = self._page_converters
 
+                if "style_map" not in _kwargs and self._style_map is not None:
+                    _kwargs["style_map"] = self._style_map
+
                 # If we hit an error log it and keep trying
                 try:
                     res = converter.convert(local_path, **_kwargs)

tests/test_markitdown.py

@@ -51,6 +51,17 @@
     "AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation",
 ]
 
+DOCX_COMMENT_TEST_STRINGS = [
+    "314b0a30-5b04-470b-b9f7-eed2c2bec74a",
+    "49e168b7-d2ae-407f-a055-2167576f39a1",
+    "## d666f1f7-46cb-42bd-9a39-9a39cf2a509f",
+    "# Abstract",
+    "# Introduction",
+    "AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation",
+    "This is a test comment. 12df-321a",
+    "Yet another comment in the doc. 55yiyi-asd09",
+]
+
 PPTX_TEST_STRINGS = [
     "2cdda5c8-e50e-4db4-b5f0-9722a649f455",
     "04191ea8-5c73-4215-a1d3-1cfb43aaaf12",
@@ -139,6 +150,24 @@ def test_markitdown_local() -> None:
         text_content = result.text_content.replace("\\", "")
         assert test_string in text_content
 
+    # Test DOCX processing, with comments
+    result = markitdown.convert(
+        os.path.join(TEST_FILES_DIR, "test_with_comment.docx"),
+        style_map="comment-reference => ",
+    )
+    for test_string in DOCX_COMMENT_TEST_STRINGS:
+        text_content = result.text_content.replace("\\", "")
+        assert test_string in text_content
+
+    # Test DOCX processing, with comments and setting style_map on init
+    markitdown_with_style_map = MarkItDown(style_map="comment-reference => ")
+    result = markitdown_with_style_map.convert(
+        os.path.join(TEST_FILES_DIR, "test_with_comment.docx")
+    )
+    for test_string in DOCX_COMMENT_TEST_STRINGS:
+        text_content = result.text_content.replace("\\", "")
+        assert test_string in text_content
+
     # Test PPTX processing
     result = markitdown.convert(os.path.join(TEST_FILES_DIR, "test.pptx"))
     for test_string in PPTX_TEST_STRINGS: