From 78ba53e82bace55c2b39140417c936da968f05c2 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 23 Feb 2026 21:30:02 +0100 Subject: [PATCH] refactor: consolidate media constraints via generic base classes (#175) Replace 6 near-identical media constraint classes with two generic bases (_MediaConstraint[M] and _MediaListConstraint[M]) that use ClassVar dispatch for artifact type and MimeType generics for type-safe supported_mime_types. Update get_constraint_input_type to use _artifact_type ClassVar fast path instead of __call__ introspection. --- src/celeste/constraints.py | 199 ++++++++++--------------------------- src/celeste/io.py | 11 +- 2 files changed, 62 insertions(+), 148 deletions(-) diff --git a/src/celeste/constraints.py b/src/celeste/constraints.py index c5820aee..d35d3b06 100644 --- a/src/celeste/constraints.py +++ b/src/celeste/constraints.py @@ -3,14 +3,13 @@ import math import re from abc import ABC, abstractmethod -from typing import Any, get_args, get_origin +from typing import Any, ClassVar, get_args, get_origin from pydantic import BaseModel, Field, computed_field from celeste.artifacts import AudioArtifact, ImageArtifact, VideoArtifact from celeste.exceptions import ConstraintViolationError -from celeste.mime_types import AudioMimeType, ImageMimeType, VideoMimeType -from celeste.types import AudioContent, ImageContent, VideoContent +from celeste.mime_types import AudioMimeType, ImageMimeType, MimeType, VideoMimeType class Constraint(BaseModel, ABC): @@ -267,193 +266,105 @@ def __call__(self, value: type[BaseModel]) -> type[BaseModel]: return value -class ImageConstraint(Constraint): - """Constraint for validating a single image artifact - validates mime_type.""" +class _MediaConstraint[M: MimeType](Constraint): + """Base for single-media-artifact constraints.""" + + _artifact_type: ClassVar[type] + _media_label: ClassVar[str] - supported_mime_types: list[ImageMimeType] | None = None - """Supported MIME types for the image.""" + supported_mime_types: list[M] | None = None - def __call__(self, value: ImageArtifact) -> ImageArtifact: - """Validate single image artifact against constraint.""" + def __call__(self, value: Any) -> Any: # noqa: ANN401 + """Validate single media artifact against constraint.""" if isinstance(value, list): - msg = "ImageConstraint requires a single ImageArtifact, not a list" + msg = f"{self.__class__.__name__} requires a single {self._artifact_type.__name__}, not a list" raise ConstraintViolationError(msg) - - if not isinstance(value, ImageArtifact): - msg = f"Must be ImageArtifact, got {type(value).__name__}" + if not isinstance(value, self._artifact_type): + msg = f"Must be {self._artifact_type.__name__}, got {type(value).__name__}" raise ConstraintViolationError(msg) - if ( self.supported_mime_types is not None - and value.mime_type not in self.supported_mime_types + and value.mime_type not in self.supported_mime_types # type: ignore[attr-defined] ): supported_values = [mt.value for mt in self.supported_mime_types] - got_value = value.mime_type.value if value.mime_type else None + got_value = value.mime_type.value if value.mime_type else None # type: ignore[attr-defined] msg = f"mime_type must be one of {supported_values}, got {got_value!r}" raise ConstraintViolationError(msg) - return value -class ImagesConstraint(Constraint): - """Constraint for validating image artifacts list - validates mime_type and count limits.""" +class _MediaListConstraint[M: MimeType](Constraint): + """Base for plural-media-artifact constraints.""" - supported_mime_types: list[ImageMimeType] | None = None - """Supported MIME types.""" + _artifact_type: ClassVar[type] + _media_label: ClassVar[str] + supported_mime_types: list[M] | None = None max_count: int | None = None - """Maximum number of images.""" - def __call__(self, value: ImageContent) -> list[ImageArtifact]: - """Validate image artifact(s) against constraint and normalize to list.""" - # Normalize: if single ImageArtifact is passed, wrap it in a list - images = value if isinstance(value, list) else [value] - - if self.max_count is not None and len(images) > self.max_count: - msg = f"Must have at most {self.max_count} image(s), got {len(images)}" + def __call__(self, value: Any) -> Any: # noqa: ANN401 + """Validate media artifact(s) against constraint and normalize to list.""" + items = value if isinstance(value, list) else [value] + if self.max_count is not None and len(items) > self.max_count: + msg = f"Must have at most {self.max_count} {self._media_label}(s), got {len(items)}" raise ConstraintViolationError(msg) - if self.supported_mime_types is not None: - for i, img in enumerate(images): - if not isinstance(img, ImageArtifact): - msg = f"Image {i + 1}: Must be ImageArtifact, got {type(img).__name__}" + label = self._media_label.capitalize() + for i, item in enumerate(items): + if not isinstance(item, self._artifact_type): + msg = f"{label} {i + 1}: Must be {self._artifact_type.__name__}, got {type(item).__name__}" raise ConstraintViolationError(msg) - if img.mime_type not in self.supported_mime_types: + if item.mime_type not in self.supported_mime_types: # type: ignore[attr-defined] supported_values = [mt.value for mt in self.supported_mime_types] - got_value = img.mime_type.value if img.mime_type else None + got_value = item.mime_type.value if item.mime_type else None # type: ignore[attr-defined] msg = ( - f"Image {i + 1}: mime_type must be one of {supported_values}, " + f"{label} {i + 1}: mime_type must be one of {supported_values}, " f"got {got_value!r}" ) raise ConstraintViolationError(msg) + return items - return images - - -class VideoConstraint(Constraint): - """Constraint for validating a single video artifact - validates mime_type.""" - - supported_mime_types: list[VideoMimeType] | None = None - """Supported MIME types for the video.""" - - def __call__(self, value: VideoArtifact) -> VideoArtifact: - """Validate single video artifact against constraint.""" - if isinstance(value, list): - msg = "VideoConstraint requires a single VideoArtifact, not a list" - raise ConstraintViolationError(msg) - if not isinstance(value, VideoArtifact): - msg = f"Must be VideoArtifact, got {type(value).__name__}" - raise ConstraintViolationError(msg) +class ImageConstraint(_MediaConstraint[ImageMimeType]): + """Constraint for validating a single image artifact - validates mime_type.""" - if ( - self.supported_mime_types is not None - and value.mime_type not in self.supported_mime_types - ): - supported_values = [mt.value for mt in self.supported_mime_types] - got_value = value.mime_type.value if value.mime_type else None - msg = f"mime_type must be one of {supported_values}, got {got_value!r}" - raise ConstraintViolationError(msg) + _artifact_type = ImageArtifact + _media_label = "image" - return value +class ImagesConstraint(_MediaListConstraint[ImageMimeType]): + """Constraint for validating image artifacts list - validates mime_type and count limits.""" -class VideosConstraint(Constraint): - """Constraint for validating video artifacts list - validates mime_type and count limits.""" + _artifact_type = ImageArtifact + _media_label = "image" - supported_mime_types: list[VideoMimeType] | None = None - """Supported MIME types.""" - max_count: int | None = None - """Maximum number of videos.""" +class VideoConstraint(_MediaConstraint[VideoMimeType]): + """Constraint for validating a single video artifact - validates mime_type.""" - def __call__(self, value: VideoContent) -> list[VideoArtifact]: - """Validate video artifact(s) against constraint and normalize to list.""" - # Normalize: if single VideoArtifact is passed, wrap it in a list - videos = value if isinstance(value, list) else [value] + _artifact_type = VideoArtifact + _media_label = "video" - if self.max_count is not None and len(videos) > self.max_count: - msg = f"Must have at most {self.max_count} video(s), got {len(videos)}" - raise ConstraintViolationError(msg) - if self.supported_mime_types is not None: - for i, vid in enumerate(videos): - if not isinstance(vid, VideoArtifact): - msg = f"Video {i + 1}: Must be VideoArtifact, got {type(vid).__name__}" - raise ConstraintViolationError(msg) - if vid.mime_type not in self.supported_mime_types: - supported_values = [mt.value for mt in self.supported_mime_types] - got_value = vid.mime_type.value if vid.mime_type else None - msg = ( - f"Video {i + 1}: mime_type must be one of {supported_values}, " - f"got {got_value!r}" - ) - raise ConstraintViolationError(msg) +class VideosConstraint(_MediaListConstraint[VideoMimeType]): + """Constraint for validating video artifacts list - validates mime_type and count limits.""" - return videos + _artifact_type = VideoArtifact + _media_label = "video" -class AudioConstraint(Constraint): +class AudioConstraint(_MediaConstraint[AudioMimeType]): """Constraint for validating a single audio artifact - validates mime_type.""" - supported_mime_types: list[AudioMimeType] | None = None - """Supported MIME types for the audio.""" + _artifact_type = AudioArtifact + _media_label = "audio" - def __call__(self, value: AudioArtifact) -> AudioArtifact: - """Validate single audio artifact against constraint.""" - if isinstance(value, list): - msg = "AudioConstraint requires a single AudioArtifact, not a list" - raise ConstraintViolationError(msg) - - if not isinstance(value, AudioArtifact): - msg = f"Must be AudioArtifact, got {type(value).__name__}" - raise ConstraintViolationError(msg) - if ( - self.supported_mime_types is not None - and value.mime_type not in self.supported_mime_types - ): - supported_values = [mt.value for mt in self.supported_mime_types] - got_value = value.mime_type.value if value.mime_type else None - msg = f"mime_type must be one of {supported_values}, got {got_value!r}" - raise ConstraintViolationError(msg) - - return value - - -class AudiosConstraint(Constraint): +class AudiosConstraint(_MediaListConstraint[AudioMimeType]): """Constraint for validating audio artifacts list - validates mime_type and count limits.""" - supported_mime_types: list[AudioMimeType] | None = None - """Supported MIME types.""" - - max_count: int | None = None - """Maximum number of audios.""" - - def __call__(self, value: AudioContent) -> list[AudioArtifact]: - """Validate audio artifact(s) against constraint and normalize to list.""" - # Normalize: if single AudioArtifact is passed, wrap it in a list - audios = value if isinstance(value, list) else [value] - - if self.max_count is not None and len(audios) > self.max_count: - msg = f"Must have at most {self.max_count} audio(s), got {len(audios)}" - raise ConstraintViolationError(msg) - - if self.supported_mime_types is not None: - for i, aud in enumerate(audios): - if not isinstance(aud, AudioArtifact): - msg = f"Audio {i + 1}: Must be AudioArtifact, got {type(aud).__name__}" - raise ConstraintViolationError(msg) - if aud.mime_type not in self.supported_mime_types: - supported_values = [mt.value for mt in self.supported_mime_types] - got_value = aud.mime_type.value if aud.mime_type else None - msg = ( - f"Audio {i + 1}: mime_type must be one of {supported_values}, " - f"got {got_value!r}" - ) - raise ConstraintViolationError(msg) - - return audios + _artifact_type = AudioArtifact + _media_label = "audio" __all__ = [ diff --git a/src/celeste/io.py b/src/celeste/io.py index ab527d93..5757e37c 100644 --- a/src/celeste/io.py +++ b/src/celeste/io.py @@ -85,10 +85,7 @@ def _extract_input_type(param_type: type) -> InputType | None: def get_constraint_input_type(constraint: Constraint) -> InputType | None: - """Get InputType from constraint's __call__ signature. - - Introspects the constraint's __call__ method to find what artifact type - it accepts, then maps to InputType using INPUT_TYPE_MAPPING. + """Get InputType from constraint's _artifact_type ClassVar or __call__ signature. Args: constraint: The constraint to inspect. @@ -96,6 +93,12 @@ def get_constraint_input_type(constraint: Constraint) -> InputType | None: Returns: InputType if the constraint accepts a mapped artifact type, None otherwise. """ + # Fast path: check _artifact_type ClassVar (media constraints) + artifact_type = getattr(constraint, "_artifact_type", None) + if artifact_type is not None and artifact_type in INPUT_TYPE_MAPPING: + return INPUT_TYPE_MAPPING[artifact_type] + + # Fallback: introspect __call__ signature (Str, other constraints) annotations = inspect.get_annotations(constraint.__call__, eval_str=True) for param_type in annotations.values(): result = _extract_input_type(param_type)