Merge pull request #161 from dreadnode/fix/data-types-imports

fix: data type imports and typing

Author: briangreunke

dreadnode/data_types/image.py

@@ -4,15 +4,32 @@
 from pathlib import Path
 
 import numpy as np
-import PIL
-import PIL.Image
 
 from dreadnode.data_types.base import DataType
 
+if t.TYPE_CHECKING:
+    import numpy as np
+
 ImageDataType = t.Union[t.Any, "np.ndarray[t.Any, t.Any]"]
 ImageDataOrPathType = str | Path | bytes | ImageDataType
 
 
+def check_imports() -> None:
+    try:
+        import PIL  # type: ignore[import-not-found]
+    except ImportError as e:
+        raise ImportError(
+            "Image processing requires Pillow. Install with: pip install dreadnode[multimodal]"
+        ) from e
+
+    try:
+        import numpy as np  # type: ignore[import-not-found]
+    except ImportError as e:
+        raise ImportError(
+            "Image processing requires NumPy. Install with: pip install dreadnode[multimodal]"
+        ) from e
+
+
 class Image(DataType):
     """
     Image media type for Dreadnode logging.
@@ -45,6 +62,7 @@ def __init__(
             caption: Optional caption for the image
             format: Optional format to use when saving (png, jpg, etc.)
         """
+        check_imports()
         self._data = data
         self._mode = mode
         self._caption = caption
@@ -66,6 +84,8 @@ def _process_image_data(self) -> tuple[bytes, str, str | None, int | None, int |
         Returns:
             A tuple of (image_bytes, image_format, mode, width, height)
         """
+        import numpy as np  # type: ignore[import-not-found]
+        import PIL.Image  # type: ignore[import-not-found]
 
         if isinstance(self._data, str | Path) and Path(self._data).exists():
             return self._process_file_path()
@@ -85,6 +105,7 @@ def _process_file_path(self) -> tuple[bytes, str, str | None, int | None, int |
         Returns:
             A tuple of (image_bytes, image_format, mode, width, height)
         """
+        import PIL.Image  # type: ignore[import-not-found]
 
         path_str = str(self._data)
         image_bytes = Path(path_str).read_bytes()
@@ -102,6 +123,7 @@ def _process_pil_image(self) -> tuple[bytes, str, str | None, int | None, int |
         Returns:
             A tuple of (image_bytes, image_format, mode, width, height)
         """
+        import PIL.Image  # type: ignore[import-not-found]
 
         if not isinstance(self._data, PIL.Image.Image):
             raise TypeError(f"Expected PIL.Image, got {type(self._data)}")
@@ -139,6 +161,8 @@ def _process_numpy_array(self) -> tuple[bytes, str, str | None, int | None, int
         Returns:
             A tuple of (image_bytes, image_format, mode, width, height)
         """
+        import numpy as np  # type: ignore[import-not-found]
+        import PIL.Image  # type: ignore[import-not-found]
 
         buffer = io.BytesIO()
         image_format = self._format or "png"
@@ -168,6 +192,7 @@ def _process_raw_bytes(self) -> tuple[bytes, str, str | None, int | None, int |
         Returns:
             A tuple of (image_bytes, image_format, mode, width, height)
         """
+        import PIL.Image  # type: ignore[import-not-found]
 
         if not isinstance(self._data, bytes):
             raise TypeError(f"Expected bytes, got {type(self._data)}")
@@ -192,6 +217,7 @@ def _process_base64_string(self) -> tuple[bytes, str, str | None, int | None, in
         Returns:
             A tuple of (image_bytes, image_format, mode, width, height)
         """
+        import PIL.Image  # type: ignore[import-not-found]
 
         if not isinstance(self._data, str):
             raise TypeError(f"Expected str, got {type(self._data)}")
@@ -228,6 +254,8 @@ def _generate_metadata(
         self, image_format: str, mode: str | None, width: int | None, height: int | None
     ) -> dict[str, str | int | None]:
         """Generate metadata for the image."""
+        import numpy as np  # type: ignore[import-not-found]
+        import PIL.Image  # type: ignore[import-not-found]
 
         metadata: dict[str, str | int | None] = {
             "extension": image_format.lower(),
@@ -286,6 +314,7 @@ def _ensure_valid_image_array(
         self, array: "np.ndarray[t.Any, np.dtype[t.Any]]"
     ) -> "np.ndarray[t.Any, np.dtype[t.Any]]":
         """Convert numpy array to a format suitable for PIL."""
+        import numpy as np  # type: ignore[import-not-found]
 
         grayscale_dim = 2
         rgb_dim = 3

dreadnode/data_types/table.py

@@ -3,16 +3,33 @@
 from pathlib import Path
 from typing import ClassVar
 
-import numpy as np
-import pandas as pd
-
 from dreadnode.data_types.base import DataType
 
+if t.TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
+
 TableDataType = t.Union[
     "pd.DataFrame", dict[t.Any, t.Any], list[t.Any], str, Path, "np.ndarray[t.Any, t.Any]"
 ]
 
 
+def check_imports() -> None:
+    try:
+        import pandas as pd  # type: ignore[import-not-found]
+    except ImportError as e:
+        raise ImportError(
+            "Image processing requires Pandas. Install with: pip install dreadnode[multimodal]"
+        ) from e
+
+    try:
+        import numpy as np  # type: ignore[import-not-found]
+    except ImportError as e:
+        raise ImportError(
+            "Image processing requires NumPy. Install with: pip install dreadnode[multimodal]"
+        ) from e
+
+
 class Table(DataType):
     """
     Table data type for Dreadnode logging.
@@ -47,6 +64,7 @@ def __init__(
             format: Optional format to use when saving (csv, parquet, json)
             index: Include index in the output
         """
+        check_imports()
         self._data = data
         self._caption = caption
         self._format = format or "csv"  # Default to CSV
@@ -77,6 +95,8 @@ def _to_dataframe(self) -> "pd.DataFrame":
         Returns:
             A pandas DataFrame representation of the input data
         """
+        import numpy as np  # type: ignore[import-not-found]
+        import pandas as pd  # type: ignore[import-not-found]
 
         if isinstance(self._data, pd.DataFrame):
             return self._data
@@ -131,6 +151,8 @@ def _generate_metadata(self, data_frame: "pd.DataFrame") -> dict[str, t.Any]:
         Returns:
             A dictionary of metadata
         """
+        import numpy as np  # type: ignore[import-not-found]
+        import pandas as pd  # type: ignore[import-not-found]
 
         metadata = {
             "extension": self._format,

pyproject.toml

@@ -109,7 +109,11 @@ python_version = "3.10"
 exclude = "tests"
 
 [[tool.mypy.overrides]]
-module = ["dreadnode.scorers.*", "dreadnode.transforms.*"]
+module = [
+    "dreadnode.data_types.*",
+    "dreadnode.scorers.*",
+    "dreadnode.transforms.*",
+]
 disable_error_code = ["unused-ignore", "import-untyped"]
 
 [tool.ty.environment]