Source code for labapi.entry.entries.attachment

"""Attachment Entry Module."""

from __future__ import annotations

import shutil
import warnings
from datetime import datetime
from email.message import Message
from io import BytesIO
from mimetypes import guess_extension
from pathlib import PurePosixPath
from tempfile import TemporaryFile
from typing import IO, TYPE_CHECKING
from urllib.parse import unquote, urlsplit

from typing_extensions import override

from labapi.entry.attachment import Attachment
from labapi.exceptions import ApiError

from .base import Entry

if TYPE_CHECKING:
    from labapi.user import User


def _make_backing_io(use_tempfile: bool) -> IO[bytes]:
    return TemporaryFile() if use_tempfile else BytesIO()


def _s3_filename_from_url(url: str) -> str | None:
    """Return a filename from a redirected Amazon S3 object URL, if valid."""
    parsed_url = urlsplit(url)
    if not parsed_url.hostname or not parsed_url.hostname.endswith(".amazonaws.com"):
        return None

    path = unquote(parsed_url.path)
    if not path or path.endswith("/"):
        return None

    filename = PurePosixPath(path).name
    if not filename.strip() or filename in {".", ".."}:
        return None

    return filename


[docs] class AttachmentEntry(Entry[Attachment], part_type="Attachment"): """Represents an attachment entry on a LabArchives page. This class handles the retrieval and updating of file attachments, providing access to the attachment's content, filename, and caption. """
[docs] def __init__( self, eid: str, caption: str, user: User, *, created_at: datetime | None = None, updated_at: datetime | None = None, version: int | None = None, ): """Initialize an attachment entry. :param eid: The unique ID of the entry. :param caption: The caption associated with the attachment. :param user: The authenticated user. """ super().__init__( eid, caption, user, created_at=created_at, updated_at=updated_at, version=version, ) self._filedata: Attachment | None = None # Filename reported by the page listing (or supplied at upload time). # This is more stable than a generated download filename. self._filename: str | None = None self._mime_type: str | None = None
def _ensure_attachment(self, use_tempfile: bool) -> None: if self._filedata is None or self._filedata.closed: with self._user.client.stream_api_get( "entries/entry_attachment", uid=self._user.id, eid=self.id ) as attachment_stream: headers = attachment_stream.headers msg = Message() msg["Content-Type"] = ( headers.get("Content-Type") or "application/octet-stream" ) content_disposition = headers.get("Content-Disposition") if content_disposition is not None: msg["Content-Disposition"] = content_disposition mime_type = msg.get_content_type() filename = self._filename or msg.get_filename() if filename is not None and not filename.strip(): filename = None if filename is None and attachment_stream.response.history: filename = _s3_filename_from_url(attachment_stream.url) if filename is None: extension = guess_extension(mime_type, strict=False) or ".bin" filename = f"{self.id}{extension}" warnings.warn( "Could not determine filename from API response headers " "or redirected S3 URL; using attachment entry EID " f"{self.id!r} with MIME-derived extension {extension!r} " f"as filename {filename!r}.", RuntimeWarning, stacklevel=2, ) output = _make_backing_io(use_tempfile) for chunk in attachment_stream: output.write(chunk) self._filedata = Attachment(output, mime_type, filename, self._data)
[docs] def get_attachment(self, use_tempfile: bool = False) -> Attachment: """Return the attachment payload as an independent stream copy. The attachment data is fetched from the LabArchives API on first call and cached. :param use_tempfile: If True, the attachment data will be stored in a temporary file; otherwise, in an in-memory BytesIO object. Defaults to False. :returns: An :class:`~labapi.entry.attachment.Attachment` object containing the file data and metadata. """ self._ensure_attachment(use_tempfile) assert self._filedata is not None output = _make_backing_io(use_tempfile) # Return an independent copy so each caller gets isolated read/seek/close state # while still sharing a single downloaded backing attachment in the cache. self._filedata.seek(0) shutil.copyfileobj(self._filedata, output) output.seek(0) return Attachment( output, self._filedata.mime_type, self._filedata.filename, self._filedata.caption, )
def _release_attachment_cache(self) -> None: """Close and discard the cached attachment payload.""" if self._filedata is not None: self._filedata.close() self._filedata = None @property @override def content(self) -> Attachment: """Return the attachment content. This property retrieves the attachment data, caching it for subsequent access. :returns: The attachment object. """ return self.get_attachment() @content.setter @override def content(self, value: Attachment): """Set the attachment content. This operation updates the attachment in LabArchives via an API call and invalidates any previously cached attachment data. :param value: The new attachment object to upload. """ # NOTE: this implicitly invalidates all previous Attachments # NOTE: if every time content is called we give a new copy anyways that's fine # (see get_attachment()) if value.seekable(): value.seek(0) try: self._user.api_post( "entries/update_attachment", value._backing, # pyright: ignore[reportPrivateUsage, reportArgumentType] filename=value.filename, caption=value.caption, eid=self.id, change_description="File updated via API", ) except ApiError as exc: if exc.error_code != 4999: raise raise ApiError( "Attachment update failed for " f"entry {self.id!r} with filename {value.filename!r}. " f"LabArchives returned {exc}. " "Reload the page or entry, revalidate the attachment metadata, " "and retry with a fresh Attachment object.", exc.error_code, ) from exc self._data = value.caption self._filename = value.filename or None if self._filedata: self._filedata.close() self._filedata = None @property def caption(self) -> str: """Return the attachment caption. :returns: The caption string. """ return self._data