Ensure TemporaryFiles are always closed after use

Author: brianhelba

isic/ingest/models/accession.py

@@ -1,24 +1,25 @@
-from collections.abc import Callable
+from collections.abc import Callable, Generator
+import contextlib
 from copy import deepcopy
 from dataclasses import dataclass
 import io
 import logging
 from mimetypes import guess_type
-from pathlib import Path, PurePosixPath
+from pathlib import PurePosixPath
 import tempfile
-from typing import Literal, TypeVar
+from typing import IO, Literal, TypeVar
 from uuid import uuid4
 
 from django.contrib.auth.models import User
 from django.contrib.postgres.constraints import ExclusionConstraint
 from django.core.exceptions import ValidationError
-from django.core.files import File
 from django.core.files.storage import storages
 from django.core.files.uploadedfile import InMemoryUploadedFile
 from django.db import models, transaction
 from django.db.models import Deferrable, FileField, FloatField, IntegerField, Transform
 from django.db.models.constraints import CheckConstraint, UniqueConstraint
 from django.db.models.fields import Field
+from django.db.models.fields.files import FieldFile
 from django.db.models.functions import Cast, Round
 from django.db.models.query_utils import Q
 from isic_metadata.fields import ImageTypeEnum
@@ -42,6 +43,12 @@
 
 logger = logging.getLogger(__name__)
 
+# Set a larger max size, to accommodate confocal images
+# This uses ~1.1GB of memory
+PIL.Image.MAX_IMAGE_PIXELS = 20_000 * 20_000 * 3
+
+gdal.UseExceptions()
+
 # Set the GDAL raster block cache to a maximum of 128MB. This is a value that
 # reduces memory usage without noticeably impacting performance for the typical mosaic.
 gdal.SetCacheMax(128 * 1024**2)
@@ -205,15 +212,6 @@ class AccessionStatus(models.TextChoices):
     SUCCEEDED = "succeeded", "Succeeded"
 
 
-@dataclass
-class AccessionBlob:
-    blob: File
-    blob_size: int
-    height: int
-    width: int
-    is_cog: bool
-
-
 def sponsored_blob_storage():
     return storages["sponsored"]
 
@@ -500,131 +498,14 @@ def metadata(self):
                 ret[field.name] = getattr(self, field.name)
         return ret
 
-    @staticmethod
-    def metadata_keys():
-        return [
-            field.name for field in Accession._meta.fields if hasattr(AccessionMetadata, field.name)
-        ]
-
-    @staticmethod
-    def is_color(img: PIL.Image.Image) -> bool:
-        return img.mode in {"RGB", "RGBA"}
-
-    @staticmethod
-    def meets_cog_threshold(img: PIL.Image.Image) -> bool:
-        return img.height * img.width > IMAGE_COG_THRESHOLD
-
     def get_diagnosis_display(self) -> str:
         diagnoses = [self.metadata.get(f"diagnosis_{i}") for i in range(1, 6)]
         if any(diagnoses):
             return [d for d in diagnoses if d is not None][-1]
         else:
             return ""
 
-    def _generate_blob(self, img: PIL.Image.Image) -> AccessionBlob:
-        # Explicitly load the image, so any decoding errors can be caught
-        try:
-            img.load()
-        except OSError as e:
-            if "image file is truncated" in str(e):
-                raise InvalidBlobError("Blob appears truncated.") from e
-
-            # Any other errors are not expected, so re-raise them natively
-            raise
-
-        # rotate the image bytes according to the orientation tag, stripping it in the process
-        PIL.ImageOps.exif_transpose(img, in_place=True)
-
-        # Strip any alpha channel
-        if self.is_color(img):
-            img = img.convert("RGB")
-
-        with tempfile.SpooledTemporaryFile() as stripped_blob_stream:
-            output_format = "PNG" if img.format == "PNG" and not self.is_color(img) else "JPEG"
-            img.save(stripped_blob_stream, format=output_format)
-
-            stripped_blob_stream.seek(0, io.SEEK_END)
-            stripped_blob_size = stripped_blob_stream.tell()
-            stripped_blob_stream.seek(0)
-
-            blob_name = f"{uuid4()}.{'png' if output_format == 'PNG' else 'jpg'}"
-            accession_blob = AccessionBlob(
-                blob=InMemoryUploadedFile(
-                    file=stripped_blob_stream,
-                    field_name=None,
-                    name=blob_name,
-                    content_type=f"image/{output_format.lower()}",
-                    size=stripped_blob_size,
-                    charset=None,
-                ),
-                blob_size=stripped_blob_size,
-                height=img.height,
-                width=img.width,
-                is_cog=False,
-            )
-            self.blob = accession_blob.blob
-            Accession._meta.get_field("blob").pre_save(self, add=False)
-
-        return accession_blob
-
-    def _generate_blob_as_cog(self, img: PIL.Image.Image) -> AccessionBlob:
-        with (
-            field_file_to_local_path(self.original_blob) as original_blob_path,
-            tempfile.NamedTemporaryFile(delete=False) as cog_temp_file,
-        ):
-            gdal.UseExceptions()
-
-            src_ds = gdal.Open(original_blob_path)
-
-            gdal.Translate(
-                cog_temp_file.name,
-                src_ds,
-                options=gdal.TranslateOptions(
-                    format="COG",
-                    # rescale unsigned 16-bit png band to 8-bit
-                    outputType=gdal.GDT_Byte,
-                    scaleParams=[[0, 2**16 - 1, 0, 2**8 - 1]],
-                    creationOptions={
-                        "BLOCKSIZE": 256,
-                        "COMPRESS": "DEFLATE",
-                        "PREDICTOR": "YES",
-                        "LEVEL": "9",
-                        "BIGTIFF": "IF_SAFER",
-                        # Strip EXIF metadata
-                        "COPY_SRC_MDD": "NO",
-                    },
-                    resampleAlg=gdal.GRA_Lanczos,
-                ),
-            )
-
-            # necessary to close the src_ds (https://gis.stackexchange.com/a/80370)
-            del src_ds
-
-        blob_size = Path(cog_temp_file.name).stat().st_size
-        with Path(cog_temp_file.name).open("rb") as cog_stream:
-            blob_name = f"{uuid4()}.tif"
-            accession_blob = AccessionBlob(
-                blob=InMemoryUploadedFile(
-                    file=cog_stream,
-                    field_name=None,
-                    name=blob_name,
-                    content_type="image/tiff",
-                    size=blob_size,
-                    charset=None,
-                ),
-                blob_size=blob_size,
-                height=img.height,
-                width=img.width,
-                is_cog=True,
-            )
-            self.blob = accession_blob.blob
-            Accession._meta.get_field("blob").pre_save(self, add=False)
-
-        Path(cog_temp_file.name).unlink()
-
-        return accession_blob
-
-    def generate_blob(self):
+    def generate_blob(self):  # noqa: PLR0915
         """
         Generate `blob` and set related attributes.
 
@@ -634,34 +515,57 @@ def generate_blob(self):
         try:
             with self.original_blob.open("rb") as original_blob_stream:
                 blob_mime_type = guess_mime_type(original_blob_stream, self.original_blob_name)
-            blob_major_mime_type = blob_mime_type.partition("/")[0]
-            if blob_major_mime_type != "image":
-                raise InvalidBlobError(  # noqa: TRY301
-                    f'Blob has a non-image MIME type: "{blob_mime_type}"'
-                )
-
-            # Set a larger max size, to accommodate confocal images
-            # This uses ~1.1GB of memory
-            PIL.Image.MAX_IMAGE_PIXELS = 20_000 * 20_000 * 3
-            try:
-                img = PIL.Image.open(original_blob_stream)
-            except PIL.Image.UnidentifiedImageError as e:
-                raise InvalidBlobError("Blob cannot be recognized by PIL.") from e
-
-            if self.meets_cog_threshold(img):
-                if self.is_color(img):
-                    raise InvalidBlobError("Blob is too large to be stored.")  # noqa: TRY301
+                if blob_mime_type.partition("/")[0] != "image":
+                    raise InvalidBlobError(  # noqa: TRY301
+                        f'Blob has a non-image MIME type: "{blob_mime_type}"'
+                    )
 
-                accession_blob = self._generate_blob_as_cog(img)
-            else:
-                accession_blob = self._generate_blob(img)
+                try:
+                    img = PIL.Image.open(original_blob_stream)
+                except PIL.Image.UnidentifiedImageError as e:
+                    raise InvalidBlobError("Blob cannot be recognized by PIL.") from e
+                self.height = img.height
+                self.width = img.width
+
+                is_rcm = not self.is_color(img)
+
+                if is_rcm:
+                    self.is_cog = self.meets_cog_threshold(img)
+                    if self.is_cog:
+                        converter = self._convert_blog_to_cog
+                        converted_blob_type = "image/tiff"
+                        converted_blob_extension = "tif"
+                    else:
+                        converter = self._convert_blob_to_png
+                        converted_blob_type = "image/png"
+                        converted_blob_extension = "png"
+                else:
+                    if self.meets_cog_threshold(img):
+                        raise InvalidBlobError("Blob is too large to be stored.")  # noqa: TRY301
+                    self.is_cog = False
+                    converter = self._convert_blob_to_rgb
+                    converted_blob_type = "image/jpeg"
+                    converted_blob_extension = "jpg"
+
+                with converter(self.original_blob) as converted_blob_stream:
+                    converted_blob_name = f"{uuid4()}.{converted_blob_extension}"
+
+                    converted_blob_stream.seek(0, io.SEEK_END)
+                    converted_blob_size = converted_blob_stream.tell()
+                    self.blob_size = converted_blob_size
+
+                    self.blob = InMemoryUploadedFile(
+                        file=converted_blob_stream,
+                        field_name=None,
+                        name=converted_blob_name,
+                        content_type=converted_blob_type,
+                        size=converted_blob_size,
+                        charset=None,
+                    )
 
-            self.blob_size = accession_blob.blob_size
-            self.height = accession_blob.height
-            self.width = accession_blob.width
-            self.is_cog = accession_blob.is_cog
+                    Accession._meta.get_field("blob").pre_save(self, add=False)
 
-            self.save(update_fields=["blob", "blob_size", "height", "width", "is_cog"])
+                    self.save(update_fields=["blob", "blob_size", "height", "width", "is_cog"])
 
             self.generate_thumbnail()
         except InvalidBlobError:
@@ -752,10 +656,6 @@ def from_blob(cls, blob: Blob):
             ),
         )
 
-    @staticmethod
-    def _age_approx(age: int) -> int:
-        return int(round(age / 5.0) * 5)
-
     @property
     def age_approx(self) -> int | None:
         return self._age_approx(self.metadata["age"]) if "age" in self.metadata else None
@@ -932,3 +832,88 @@ def full_clean(self, *args, **kwargs):
             raise Exception("unstructured_metadata is required")
 
         super().full_clean(*args, **kwargs)
+
+    @staticmethod
+    def _age_approx(age: int) -> int:
+        return int(round(age / 5.0) * 5)
+
+    @staticmethod
+    def metadata_keys():
+        return [
+            field.name for field in Accession._meta.fields if hasattr(AccessionMetadata, field.name)
+        ]
+
+    @staticmethod
+    def is_color(img: PIL.Image.Image) -> bool:
+        return img.mode in {"RGB", "RGBA"}
+
+    @staticmethod
+    def meets_cog_threshold(img: PIL.Image.Image) -> bool:
+        return img.height * img.width > IMAGE_COG_THRESHOLD
+
+    @staticmethod
+    def _ensure_pil_image(img: PIL.Image.Image) -> None:
+        """Explicitly load a PIL image's pixel data, so any decoding errors can be caught."""
+        try:
+            img.load()
+        except OSError as e:
+            if "image file is truncated" in str(e):
+                raise InvalidBlobError("Blob appears truncated.") from e
+            # Any other errors are not expected, so re-raise them natively
+            raise
+
+    @staticmethod
+    @contextlib.contextmanager
+    def _convert_blob_to_rgb(blob: FieldFile) -> Generator[IO[bytes]]:
+        img = PIL.Image.open(blob)
+        Accession._ensure_pil_image(img)
+
+        # rotate the image bytes according to the orientation tag, stripping it in the process
+        PIL.ImageOps.exif_transpose(img, in_place=True)
+
+        # Strip any alpha channel
+        img = img.convert("RGB")
+
+        with tempfile.SpooledTemporaryFile() as converted_blob_stream:
+            img.save(converted_blob_stream, format="JPEG")
+            yield converted_blob_stream
+
+    @staticmethod
+    @contextlib.contextmanager
+    def _convert_blob_to_png(blob: FieldFile) -> Generator[IO[bytes]]:
+        img = PIL.Image.open(blob)
+        Accession._ensure_pil_image(img)
+
+        with tempfile.SpooledTemporaryFile() as converted_blob_stream:
+            img.save(converted_blob_stream, format="PNG")
+            yield converted_blob_stream
+
+    @staticmethod
+    @contextlib.contextmanager
+    def _convert_blog_to_cog(blob: FieldFile) -> Generator[IO[bytes]]:
+        with tempfile.NamedTemporaryFile() as converted_blob_stream:
+            with (
+                field_file_to_local_path(blob) as blob_path,
+                gdal.Open(blob_path) as src_dataset,
+            ):
+                gdal.Translate(
+                    converted_blob_stream.name,
+                    src_dataset,
+                    options=gdal.TranslateOptions(
+                        format="COG",
+                        # rescale unsigned 16-bit png band to 8-bit
+                        outputType=gdal.GDT_Byte,
+                        scaleParams=[[0, 2**16 - 1, 0, 2**8 - 1]],
+                        creationOptions={
+                            "BLOCKSIZE": 256,
+                            "COMPRESS": "DEFLATE",
+                            "PREDICTOR": "YES",
+                            "LEVEL": "9",
+                            "BIGTIFF": "IF_SAFER",
+                            # Strip EXIF metadata
+                            "COPY_SRC_MDD": "NO",
+                        },
+                        resampleAlg=gdal.GRA_Lanczos,
+                    ),
+                )
+            yield converted_blob_stream