fix: add handling for str filenames in save/load methods (#205)

* fix: add handling for str filenames in save/load methods

Signed-off-by: Musashi Hinck <[email protected]>

* fix: add handling for str filenames in save/load methods

Signed-off-by: Musashi Hinck <[email protected]>

---------

Signed-off-by: Musashi Hinck <[email protected]>
Co-authored-by: Musashi Hinck <[email protected]>

Author: muhark

docling_core/types/doc/document.py

@@ -2410,12 +2410,14 @@ def export_to_element_tree(self) -> str:
 
     def save_as_json(
         self,
-        filename: Path,
+        filename: Union[str, Path],
         artifacts_dir: Optional[Path] = None,
         image_mode: ImageRefMode = ImageRefMode.EMBEDDED,
         indent: int = 2,
     ):
         """Save as json."""
+        if isinstance(filename, str):
+            filename = Path(filename)
         artifacts_dir, reference_path = self._get_output_paths(filename, artifacts_dir)
 
         if image_mode == ImageRefMode.REFERENCED:
@@ -2430,7 +2432,7 @@ def save_as_json(
             json.dump(out, fw, indent=indent)
 
     @classmethod
-    def load_from_json(cls, filename: Path) -> "DoclingDocument":
+    def load_from_json(cls, filename: Union[str, Path]) -> "DoclingDocument":
         """load_from_json.
 
         :param filename: The filename to load a saved DoclingDocument from a .json.
@@ -2440,17 +2442,21 @@ def load_from_json(cls, filename: Path) -> "DoclingDocument":
         :rtype: DoclingDocument
 
         """
+        if isinstance(filename, str):
+            filename = Path(filename)
         with open(filename, "r", encoding="utf-8") as f:
             return cls.model_validate_json(f.read())
 
     def save_as_yaml(
         self,
-        filename: Path,
+        filename: Union[str, Path],
         artifacts_dir: Optional[Path] = None,
         image_mode: ImageRefMode = ImageRefMode.EMBEDDED,
         default_flow_style: bool = False,
     ):
         """Save as yaml."""
+        if isinstance(filename, str):
+            filename = Path(filename)
         artifacts_dir, reference_path = self._get_output_paths(filename, artifacts_dir)
 
         if image_mode == ImageRefMode.REFERENCED:
@@ -2465,7 +2471,7 @@ def save_as_yaml(
             yaml.dump(out, fw, default_flow_style=default_flow_style)
 
     @classmethod
-    def load_from_yaml(cls, filename: Path) -> "DoclingDocument":
+    def load_from_yaml(cls, filename: Union[str, Path]) -> "DoclingDocument":
         """load_from_yaml.
 
         Args:
@@ -2474,6 +2480,8 @@ def load_from_yaml(cls, filename: Path) -> "DoclingDocument":
         Returns:
             DoclingDocument: the loaded DoclingDocument
         """
+        if isinstance(filename, str):
+            filename = Path(filename)
         with open(filename, encoding="utf-8") as f:
             data = yaml.load(f, Loader=yaml.FullLoader)
         return DoclingDocument.model_validate(data)
@@ -2491,7 +2499,7 @@ def export_to_dict(
 
     def save_as_markdown(
         self,
-        filename: Path,
+        filename: Union[str, Path],
         artifacts_dir: Optional[Path] = None,
         delim: str = "\n\n",
         from_element: int = 0,
@@ -2507,6 +2515,8 @@ def save_as_markdown(
         included_content_layers: set[ContentLayer] = DEFAULT_CONTENT_LAYERS,
     ):
         """Save to markdown."""
+        if isinstance(filename, str):
+            filename = Path(filename)
         artifacts_dir, reference_path = self._get_output_paths(filename, artifacts_dir)
 
         if image_mode == ImageRefMode.REFERENCED:
@@ -2634,7 +2644,7 @@ def export_to_text(  # noqa: C901
 
     def save_as_html(
         self,
-        filename: Path,
+        filename: Union[str, Path],
         artifacts_dir: Optional[Path] = None,
         from_element: int = 0,
         to_element: int = sys.maxsize,
@@ -2647,6 +2657,8 @@ def save_as_html(
         included_content_layers: set[ContentLayer] = DEFAULT_CONTENT_LAYERS,
     ):
         """Save to HTML."""
+        if isinstance(filename, str):
+            filename = Path(filename)
         artifacts_dir, reference_path = self._get_output_paths(filename, artifacts_dir)
 
         if image_mode == ImageRefMode.REFERENCED:
@@ -2672,8 +2684,10 @@ def save_as_html(
             fw.write(html_out)
 
     def _get_output_paths(
-        self, filename: Path, artifacts_dir: Optional[Path] = None
+        self, filename: Union[str, Path], artifacts_dir: Optional[Path] = None
     ) -> Tuple[Path, Optional[Path]]:
+        if isinstance(filename, str):
+            filename = Path(filename)
         if artifacts_dir is None:
             # Remove the extension and add '_pictures'
             artifacts_dir = filename.with_suffix("")
@@ -3455,7 +3469,7 @@ def save_as_document_tokens(self, *args, **kwargs):
 
     def save_as_doctags(
         self,
-        filename: Path,
+        filename: Union[str, Path],
         delim: str = "",
         from_element: int = 0,
         to_element: int = sys.maxsize,
@@ -3470,6 +3484,8 @@ def save_as_doctags(
         add_table_cell_text: bool = True,
     ):
         r"""Save the document content to DocTags format."""
+        if isinstance(filename, str):
+            filename = Path(filename)
         out = self.export_to_document_tokens(
             delim=delim,
             from_element=from_element,

docling_core/types/doc/page.py

@@ -546,7 +546,7 @@ def export_to_dict(self) -> Dict:
 
     def save_as_json(
         self,
-        filename: Path,
+        filename: Union[str, Path],
         indent: int = 2,
     ):
         """Save the page data as a JSON file.
@@ -555,12 +555,14 @@ def save_as_json(
             filename: Path to save the JSON file
             indent: Indentation level for JSON formatting
         """
+        if isinstance(filename, str):
+            filename = Path(filename)
         out = self.export_to_dict()
         with open(filename, "w", encoding="utf-8") as fw:
             json.dump(out, fw, indent=indent)
 
     @classmethod
-    def load_from_json(cls, filename: Path) -> "SegmentedPdfPage":
+    def load_from_json(cls, filename: Union[str, Path]) -> "SegmentedPdfPage":
         """Load page data from a JSON file.
 
         Args:
@@ -569,6 +571,8 @@ def load_from_json(cls, filename: Path) -> "SegmentedPdfPage":
         Returns:
             Instantiated SegmentedPdfPage object
         """
+        if isinstance(filename, str):
+            filename = Path(filename)
         with open(filename, "r", encoding="utf-8") as f:
             return cls.model_validate_json(f.read())
 
@@ -1155,19 +1159,21 @@ def export_to_dict(self, mode: str = "json") -> Dict:
         """
         return self.model_dump(mode=mode, by_alias=True, exclude_none=True)
 
-    def save_as_json(self, filename: Path, indent: int = 2):
+    def save_as_json(self, filename: Union[str, Path], indent: int = 2):
         """Save the table of contents as a JSON file.
 
         Args:
             filename: Path to save the JSON file
             indent: Indentation level for JSON formatting
         """
+        if isinstance(filename, str):
+            filename = Path(filename)
         out = self.export_to_dict()
         with open(filename, "w", encoding="utf-8") as fw:
             json.dump(out, fw, indent=indent)
 
     @classmethod
-    def load_from_json(cls, filename: Path) -> "PdfTableOfContents":
+    def load_from_json(cls, filename: Union[str, Path]) -> "PdfTableOfContents":
         """Load table of contents from a JSON file.
 
         Args:
@@ -1176,6 +1182,8 @@ def load_from_json(cls, filename: Path) -> "PdfTableOfContents":
         Returns:
             Instantiated PdfTableOfContents object
         """
+        if isinstance(filename, str):
+            filename = Path(filename)
         with open(filename, "r", encoding="utf-8") as f:
             return cls.model_validate_json(f.read())
 
@@ -1213,19 +1221,21 @@ def export_to_dict(
         """
         return self.model_dump(mode=mode, by_alias=True, exclude_none=True)
 
-    def save_as_json(self, filename: Path, indent: int = 2):
+    def save_as_json(self, filename: Union[str, Path], indent: int = 2):
         """Save the document as a JSON file.
 
         Args:
             filename: Path to save the JSON file
             indent: Indentation level for JSON formatting
         """
+        if isinstance(filename, str):
+            filename = Path(filename)
         out = self.export_to_dict()
         with open(filename, "w", encoding="utf-8") as fw:
             json.dump(out, fw, indent=indent)
 
     @classmethod
-    def load_from_json(cls, filename: Path) -> "ParsedPdfDocument":
+    def load_from_json(cls, filename: Union[str, Path]) -> "ParsedPdfDocument":
         """Load document from a JSON file.
 
         Args:
@@ -1234,5 +1244,7 @@ def load_from_json(cls, filename: Path) -> "ParsedPdfDocument":
         Returns:
             Instantiated ParsedPdfDocument object
         """
+        if isinstance(filename, str):
+            filename = Path(filename)
         with open(filename, "r", encoding="utf-8") as f:
             return cls.model_validate_json(f.read())