refactor(file-utils): enhance file handling utilities for images and PDFs

- Breaking change: Required `datatype` in data url to detect supported file formats
- Introduced a new `io_utils.py` module for determining MIME types and reading bytes from various sources.
- Refactored `load_source` in `file_utils.py` to utilize the new `source_file_type` function for improved file type detection.
- Updated `ImageSource` and `PdfSource` classes to leverage the `read_bytes` function for loading data.
- Removed the deprecated `load_image` and `load_pdf` functions, streamlining the codebase.
- Enhanced tests for `PdfSource` and `ImageSource` to validate loading from paths and data URLs, ensuring robust error handling.

Author: danyalxahid-askui

src/askui/utils/file_utils.py

@@ -1,33 +1,49 @@
 from pathlib import Path
 from typing import Union
 
-from filetype import guess  # type: ignore[import-untyped]
-from PIL import Image
+from PIL import Image as PILImage
 
 from askui.utils.image_utils import ImageSource
+from askui.utils.io_utils import source_file_type
 from askui.utils.pdf_utils import PdfSource
 
+# to avoid circular imports from image_utils and pdf_utils on read_bytes
 Source = Union[ImageSource, PdfSource]
 
+ALLOWED_IMAGE_TYPES = [
+    "image/png",
+    "image/jpeg",
+    "image/gif",
+    "image/webp",
+]
 
-def load_source(source: Union[str, Path, Image.Image]) -> Source:
-    """Load a source and return appropriate Source object based on file type."""
+PDF_TYPE = "application/pdf"
 
-    if isinstance(source, Image.Image):
-        return ImageSource(source)
+ALLOWED_MIMETYPES = [PDF_TYPE] + ALLOWED_IMAGE_TYPES
 
-    filepath = Path(source)
-    if not filepath.is_file():
-        msg = f"No such file or directory: '{source}'"
-        raise FileNotFoundError(msg)
 
-    kind = guess(str(filepath))
-    if kind and kind.mime == "application/pdf":
-        return PdfSource(source)
-    if kind and kind.mime.startswith("image/"):
+def load_source(source: Union[str, Path, PILImage.Image]) -> Source:
+    """Load a source and return it as an ImageSource or PdfSource.
+
+    Args:
+        source (Union[str, Path]): The source to load.
+
+    Returns:
+        Source: The loaded source as an ImageSource or PdfSource.
+
+    Raises:
+        ValueError: If the source is not a valid image or PDF file.
+    """
+    if isinstance(source, PILImage.Image):
+        return ImageSource(source)
+
+    file_type = source_file_type(source)
+    if file_type in ALLOWED_IMAGE_TYPES:
         return ImageSource(source)
-    msg = f"Unsupported file type: {filepath.suffix}"
+    if file_type == PDF_TYPE:
+        return PdfSource(source)
+    msg = f"Unsupported file type: {file_type}"
     raise ValueError(msg)
 
 
-__all__ = ["load_source", "Source"]
+__all__ = ["Source", "load_source"]

src/askui/utils/image_utils.py

@@ -11,52 +11,20 @@
 from PIL import Image as PILImage
 from pydantic import ConfigDict, RootModel, field_validator
 
+from askui.utils.io_utils import read_bytes
+
 # Regex to capture any kind of valid base64 data url (with optional media type and ;base64)
 # e.g., data:image/png;base64,... or data:;base64,... or data:,... or just ,...
 _DATA_URL_GENERIC_RE = re.compile(r"^(?:data:)?[^,]*?,(.*)$", re.DOTALL)
 
 
-def load_image(source: Union[str, Path, Image.Image]) -> Image.Image:
-    """Load and validate an image from a PIL Image, a path, or any form of base64 data URL.
-
-    Args:
-        source (Union[str, Path, Image.Image]): The image source to load from.
-            Can be a PIL Image, file path (`str` or `pathlib.Path`), or data URL.
-
-    Returns:
-        Image.Image: A valid PIL Image object.
-
-    Raises:
-        ValueError: If the input is not a valid or recognizable image.
-    """
-    if isinstance(source, Image.Image):
-        return source
-
-    if isinstance(source, Path) or (not source.startswith(("data:", ","))):
-        try:
-            return Image.open(source)
-        except (OSError, FileNotFoundError, UnidentifiedImageError) as e:
-            error_msg = f"Could not open image from file path: {source}"
-            raise ValueError(error_msg) from e
-
-    else:
-        match = _DATA_URL_GENERIC_RE.match(source)
-        if match:
-            try:
-                image_data = base64.b64decode(match.group(1))
-                return Image.open(io.BytesIO(image_data))
-            except (binascii.Error, UnidentifiedImageError):
-                try:
-                    return Image.open(source)
-                except (FileNotFoundError, UnidentifiedImageError) as e:
-                    error_msg = (
-                        f"Could not decode or identify image from input:"
-                        f"{source[:100]}{'...' if len(source) > 100 else ''}"
-                    )
-                    raise ValueError(error_msg) from e
-
-    error_msg = f"Unsupported image input type: {type(source)}"
-    raise ValueError(error_msg)
+def _bytes_to_image(image_bytes: bytes) -> Image.Image:
+    """Convert bytes to a PIL Image."""
+    try:
+        return Image.open(io.BytesIO(image_bytes))
+    except (FileNotFoundError, UnidentifiedImageError) as e:
+        error_msg = "Could not identify image from bytes"
+        raise ValueError(error_msg) from e
 
 
 def image_to_data_url(image: PILImage.Image) -> str:
@@ -391,8 +359,12 @@ def __init__(self, root: Img, **kwargs: dict[str, Any]) -> None:
 
     @field_validator("root", mode="before")
     @classmethod
-    def validate_root(cls, v: Any) -> PILImage.Image:
-        return load_image(v)
+    def validate_root(cls, v: Any) -> Image.Image:
+        if isinstance(v, Image.Image):
+            return v
+
+        image_bytes = read_bytes(v)
+        return _bytes_to_image(image_bytes)
 
     def to_data_url(self) -> str:
         """Convert the image to a data URL.
@@ -422,7 +394,6 @@ def to_bytes(self) -> bytes:
 
 
 __all__ = [
-    "load_image",
     "image_to_data_url",
     "data_url_to_image",
     "draw_point_on_image",

src/askui/utils/io_utils.py

@@ -0,0 +1,77 @@
+import base64
+import binascii
+import re
+from pathlib import Path
+from typing import Union
+
+from filetype import guess  # type: ignore[import-untyped]
+
+_DATA_URL_WITH_MIMETYPE_RE = re.compile(r"^data:([^;,]+)[^,]*?,(.*)$", re.DOTALL)
+
+
+def source_file_type(source: Union[str, Path]) -> str:
+    """Determines the MIME type of a source.
+
+    The source can be a file path or a data URL.
+
+    Args:
+        source (Union[str , Path]): The source to determine the type of.
+            Can be a file path (`str` or `pathlib.Path`) or a data URL.
+
+    Returns:
+        str: The MIME type of the source, or "unknown" if it cannot be determined.
+    """
+
+    # when source is a data url
+    if isinstance(source, str) and source.startswith("data:"):
+        match = _DATA_URL_WITH_MIMETYPE_RE.match(source)
+        if match and match.group(1):
+            return match.group(1)
+    else:
+        kind = guess(str(source))
+        if kind is not None and kind.mime is not None:
+            return str(kind.mime)
+
+    return "unknown"
+
+
+def read_bytes(source: Union[str, Path]) -> bytes:
+    """Read the bytes of a source.
+
+    The source can be a file path or a data URL.
+
+    Args:
+        source (Union[str, Path]): The source to read the bytes from.
+
+    Returns:
+        bytes: The content of the source as bytes.
+    """
+    # when source is a file path and not a data url
+    if isinstance(source, Path) or (
+        isinstance(source, str) and not source.startswith(("data:", ","))
+    ):
+        filepath = Path(source)
+        if not filepath.is_file():
+            err_msg = f"No such file or directory: '{source}'"
+            raise ValueError(err_msg)
+
+        return filepath.read_bytes()
+
+    # when source is a data url
+    if isinstance(source, str) and source.startswith(("data:", ",")):
+        match = _DATA_URL_WITH_MIMETYPE_RE.match(source)
+        if match:
+            try:
+                return base64.b64decode(match.group(2))
+            except binascii.Error as e:
+                error_msg = (
+                    "Could not decode base64 data from input: "
+                    f"{source[:100]}{'...' if len(source) > 100 else ''}"
+                )
+                raise ValueError(error_msg) from e
+
+    msg = f"Unsupported source type: {type(source)}"
+    raise ValueError(msg)
+
+
+__all__ = ["read_bytes"]

src/askui/utils/pdf_utils.py

@@ -3,6 +3,8 @@
 
 from pydantic import ConfigDict, RootModel, field_validator
 
+from askui.utils.io_utils import read_bytes
+
 Pdf = Union[str, Path]
 """Type of the input PDFs for `askui.VisionAgent.get()`, etc.
 
@@ -11,27 +13,6 @@
 """
 
 
-def load_pdf(source: Union[str, Path]) -> bytes:
-    """Load a PDF from a path and return its bytes.
-
-    Args:
-        source (Union[str, Path]): The PDF source to load from.
-
-    Returns:
-        bytes: The PDF content as bytes.
-
-    Raises:
-        FileNotFoundError: If the file is not found.
-        ValueError: If the file is too large.
-    """
-    filepath = Path(source)
-    if not filepath.is_file():
-        err_msg = f"No such file or directory: '{source}'"
-        raise FileNotFoundError(err_msg)
-
-    return filepath.read_bytes()
-
-
 class PdfSource(RootModel):
     """A class that represents a PDF source.
     It provides methods to convert it to different formats.
@@ -55,11 +36,10 @@ def __init__(self, root: Pdf, **kwargs: dict[str, Any]) -> None:
     @field_validator("root", mode="before")
     @classmethod
     def validate_root(cls, v: Any) -> bytes:
-        return load_pdf(v)
+        return read_bytes(v)
 
 
 __all__ = [
     "PdfSource",
     "Pdf",
-    "load_pdf",
 ]

tests/integration/utils/test_pdf_utils.py

@@ -1,22 +1,41 @@
+import base64
 import pathlib
 
 import pytest
 
-from askui.utils.pdf_utils import load_pdf
+from askui.utils.file_utils import PdfSource
 
 
 class TestLoadPdf:
     def test_load_pdf_from_path(self, path_fixtures_dummy_pdf: pathlib.Path) -> None:
         # Test loading from Path
-        loaded = load_pdf(path_fixtures_dummy_pdf)
-        assert isinstance(loaded, bytes)
-        assert len(loaded) > 0
+        loaded = PdfSource(path_fixtures_dummy_pdf)
+        assert isinstance(loaded.root, bytes)
+        assert len(loaded.root) > 0
 
         # Test loading from str path
-        loaded = load_pdf(str(path_fixtures_dummy_pdf))
-        assert isinstance(loaded, bytes)
-        assert len(loaded) > 0
+        loaded = PdfSource(str(path_fixtures_dummy_pdf))
+        assert isinstance(loaded.root, bytes)
+        assert len(loaded.root) > 0
 
     def test_load_pdf_nonexistent_file(self) -> None:
-        with pytest.raises(FileNotFoundError):
-            load_pdf("nonexistent_file.pdf")
+        with pytest.raises(ValueError):
+            PdfSource("nonexistent_file.pdf")
+
+    def test_pdf_source_from_data_url(
+        self, path_fixtures_dummy_pdf: pathlib.Path
+    ) -> None:
+        # Load test image and convert to base64
+        with pathlib.Path.open(path_fixtures_dummy_pdf, "rb") as f:
+            pdf_bytes = f.read()
+        pdf_str = base64.b64encode(pdf_bytes).decode()
+
+        # Test different base64 formats
+        formats = [
+            f"data:application/pdf;base64,{pdf_str}",
+        ]
+
+        for fmt in formats:
+            source = PdfSource(fmt)
+            assert isinstance(source.root, bytes)
+            assert len(source.root) > 0