feat: extend source resolution with streams and workdir (#79)

Signed-off-by: Michele Dolfi <[email protected]>
Signed-off-by: Panos Vagenas <[email protected]>
Signed-off-by: Michele Dolfi <[email protected]>
Co-authored-by: Panos Vagenas <[email protected]>

Author: dolfim-ibm

docling_core/types/io/__init__.py

@@ -0,0 +1,19 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+
+"""Models for io."""
+
+from io import BytesIO
+
+from pydantic import BaseModel, ConfigDict
+
+
+class DocumentStream(BaseModel):
+    """Wrapper class for a bytes stream with a filename."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    name: str
+    stream: BytesIO

docling_core/utils/file.py

@@ -7,28 +7,62 @@
 
 import importlib
 import tempfile
+from io import BytesIO
 from pathlib import Path
 from typing import Dict, Optional, Union
 
 import requests
 from pydantic import AnyHttpUrl, TypeAdapter, ValidationError
+from typing_extensions import deprecated
 
+from docling_core.types.io import DocumentStream
 
-def resolve_file_source(
-    source: Union[Path, AnyHttpUrl, str], headers: Optional[Dict[str, str]] = None
-) -> Path:
-    """Resolves the source (URL, path) of a file to a local file path.
 
-    If a URL is provided, the content is first downloaded to a temporary local file.
+def resolve_remote_filename(
+    http_url: AnyHttpUrl,
+    response_headers: Dict[str, str],
+    fallback_filename="file",
+) -> str:
+    """Resolves the filename from a remote url and its response headers.
+
+    Args:
+        source AnyHttpUrl: The source http url.
+        response_headers Dict: Headers received while fetching the remote file.
+        fallback_filename str: Filename to use in case none can be determined.
+
+    Returns:
+        str: The actual filename of the remote url.
+    """
+    fname = None
+    # try to get filename from response header
+    if cont_disp := response_headers.get("Content-Disposition"):
+        for par in cont_disp.strip().split(";"):
+            # currently only handling directive "filename" (not "*filename")
+            if (split := par.split("=")) and split[0].strip() == "filename":
+                fname = "=".join(split[1:]).strip().strip("'\"") or None
+                break
+    # otherwise, use name from URL:
+    if fname is None:
+        fname = Path(http_url.path or "").name or fallback_filename
+
+    return fname
+
+
+def resolve_source_to_stream(
+    source: Union[Path, AnyHttpUrl, str], headers: Optional[Dict[str, str]] = None
+) -> DocumentStream:
+    """Resolves the source (URL, path) of a file to a binary stream.
 
     Args:
         source (Path | AnyHttpUrl | str): The file input source. Can be a path or URL.
+        headers (Dict | None): Optional set of headers to use for fetching
+            the remote URL.
 
     Raises:
         ValueError: If source is of unexpected type.
 
     Returns:
-        Path: The local file path.
+        DocumentStream: The resolved file loaded as a stream.
     """
     try:
         http_url: AnyHttpUrl = TypeAdapter(AnyHttpUrl).validate_python(source)
@@ -44,29 +78,98 @@ def resolve_file_source(
         # fetch the page
         res = requests.get(http_url, stream=True, headers=req_headers)
         res.raise_for_status()
-        fname = None
-        # try to get filename from response header
-        if cont_disp := res.headers.get("Content-Disposition"):
-            for par in cont_disp.strip().split(";"):
-                # currently only handling directive "filename" (not "*filename")
-                if (split := par.split("=")) and split[0].strip() == "filename":
-                    fname = "=".join(split[1:]).strip().strip("'\"") or None
-                    break
-        # otherwise, use name from URL:
-        if fname is None:
-            fname = Path(http_url.path or "").name or "file"
-        local_path = Path(tempfile.mkdtemp()) / fname
-        with open(local_path, "wb") as f:
-            for chunk in res.iter_content(chunk_size=1024):  # using 1-KB chunks
-                f.write(chunk)
+        fname = resolve_remote_filename(http_url=http_url, response_headers=res.headers)
+
+        stream = BytesIO(res.content)
+        doc_stream = DocumentStream(name=fname, stream=stream)
     except ValidationError:
         try:
             local_path = TypeAdapter(Path).validate_python(source)
+            stream = BytesIO(local_path.read_bytes())
+            doc_stream = DocumentStream(name=local_path.name, stream=stream)
         except ValidationError:
             raise ValueError(f"Unexpected source type encountered: {type(source)}")
+    return doc_stream
+
+
+def _resolve_source_to_path(
+    source: Union[Path, AnyHttpUrl, str],
+    headers: Optional[Dict[str, str]] = None,
+    workdir: Optional[Path] = None,
+) -> Path:
+    doc_stream = resolve_source_to_stream(source=source, headers=headers)
+
+    # use a temporary directory if not specified
+    if workdir is None:
+        workdir = Path(tempfile.mkdtemp())
+
+    # create the parent workdir if it doesn't exist
+    workdir.mkdir(exist_ok=True, parents=True)
+
+    # save result to a local file
+    local_path = workdir / doc_stream.name
+    with local_path.open("wb") as f:
+        f.write(doc_stream.stream.read())
+
     return local_path
 
 
+def resolve_source_to_path(
+    source: Union[Path, AnyHttpUrl, str],
+    headers: Optional[Dict[str, str]] = None,
+    workdir: Optional[Path] = None,
+) -> Path:
+    """Resolves the source (URL, path) of a file to a local file path.
+
+    If a URL is provided, the content is first downloaded to a local file, located in
+      the provided workdir or in a temporary directory if no workdir provided.
+
+    Args:
+        source (Path | AnyHttpUrl | str): The file input source. Can be a path or URL.
+        headers (Dict | None): Optional set of headers to use for fetching
+            the remote URL.
+        workdir (Path | None): If set, the work directory where the file will
+            be downloaded, otherwise a temp dir will be used.
+
+    Raises:
+        ValueError: If source is of unexpected type.
+
+    Returns:
+        Path: The local file path.
+    """
+    return _resolve_source_to_path(
+        source=source,
+        headers=headers,
+        workdir=workdir,
+    )
+
+
+@deprecated("Use `resolve_source_to_path()` or `resolve_source_to_stream()`  instead")
+def resolve_file_source(
+    source: Union[Path, AnyHttpUrl, str],
+    headers: Optional[Dict[str, str]] = None,
+) -> Path:
+    """Resolves the source (URL, path) of a file to a local file path.
+
+    If a URL is provided, the content is first downloaded to a temporary local file.
+
+    Args:
+        source (Path | AnyHttpUrl | str): The file input source. Can be a path or URL.
+        headers (Dict | None): Optional set of headers to use for fetching
+            the remote URL.
+
+    Raises:
+        ValueError: If source is of unexpected type.
+
+    Returns:
+        Path: The local file path.
+    """
+    return _resolve_source_to_path(
+        source=source,
+        headers=headers,
+    )
+
+
 def relative_path(src: Path, target: Path) -> Path:
     """Compute the relative path from `src` to `target`.
 

poetry.lock

@@ -54,6 +54,7 @@ tabulate = "^0.9.0"
 pandas = "^2.1.4"
 pillow = "^10.3.0"
 pyyaml = ">=5.1,<7.0.0"
+typing-extensions = "^4.12.2"
 
 [tool.poetry.group.dev.dependencies]
 black = "^24.4.2"

pyproject.toml

@@ -10,7 +10,7 @@
 from requests import Response
 
 from docling_core.utils.alias import AliasModel
-from docling_core.utils.file import resolve_file_source
+from docling_core.utils.file import resolve_source_to_path, resolve_source_to_stream
 
 
 def test_alias_model():
@@ -51,7 +51,7 @@ class AliasModelGrandChild(AliasModelChild):
     assert obj.model_dump_json() != json.dumps(data, separators=(",", ":"))
 
 
-def test_resolve_file_source_url_wout_path(monkeypatch):
+def test_resolve_source_to_path_url_wout_path(monkeypatch):
     expected_str = "foo"
     expected_bytes = bytes(expected_str, "utf-8")
 
@@ -66,7 +66,29 @@ def get_dummy_response(*args, **kwargs):
         "requests.models.Response.iter_content",
         lambda *args, **kwargs: [expected_bytes],
     )
-    path = resolve_file_source("https://pypi.org")
+    path = resolve_source_to_path("https://pypi.org")
     with open(path) as f:
         text = f.read()
     assert text == expected_str
+
+
+def test_resolve_source_to_stream_url_wout_path(monkeypatch):
+    expected_str = "foo"
+    expected_bytes = bytes(expected_str, "utf-8")
+
+    def get_dummy_response(*args, **kwargs):
+        r = Response()
+        r.status_code = 200
+        r._content = expected_bytes
+        return r
+
+    monkeypatch.setattr("requests.get", get_dummy_response)
+    monkeypatch.setattr(
+        "requests.models.Response.iter_content",
+        lambda *args, **kwargs: [expected_bytes],
+    )
+    doc_stream = resolve_source_to_stream("https://pypi.org")
+    assert doc_stream.name == "file"
+
+    text = doc_stream.stream.read().decode("utf8")
+    assert text == expected_str