""" Data models for the DB and API. """ from dataclasses import dataclass from functools import partial from typing import Annotated, Any, Self import numpy as np from pydantic import BaseModel, Field, RootModel, StringConstraints, model_validator from .constants import FlagResolutions from .typings import ( ISRC, UPC, AssetID, AuditGroupID, AuditID, Email, FlagID, FlagText, LabelID, PermissionLevel, PermissionScope, UserID, VideoID, ) OptionalField = partial(Field, None) class User(BaseModel): """User data model.""" # Fields are not required by default so to allow partial updates. id: UserID = OptionalField(description="User ID.") name: str = OptionalField(description="User name.") nickname: str = OptionalField(description="User nickname.") email: Email = OptionalField(description="User email.") locale: str = OptionalField(description="User locale.") timezone: str = OptionalField(description="User timezone.") # Permission model (a dictionary of permissions) Permissions = RootModel[dict[PermissionScope, PermissionLevel | None]] class AuditGroupIds(BaseModel): """Audit group IDs.""" audit_group_ids: list[AuditGroupID] class AuditIds(BaseModel): """Audit IDs.""" audit_ids: list[AuditID] class NewAudit(BaseModel): """New audit request.""" label_id: LabelID = Field(..., description="Label ID.") audit_audio: bool = Field( True, description="Whether to include sound recordings in the audit." ) audit_video: bool = Field( False, description="Whether to include music videos in the audit." ) audit_art_track: bool = Field( False, description="Whether to include art tracks in the audit." ) @model_validator(mode="after") def check_at_least_one_audit_type_included(self) -> Self: if not any([self.audit_audio, self.audit_video, self.audit_art_track]): raise ValueError("At least one audit type must be included.") return self class FlagResolve(BaseModel): """Resolve flags request.""" flag_ids: list[FlagID] = Field(..., description="List of flag IDs to resolve.") resolution: str | None = Field( None, description="Resolution (e.g. RESOLUTION_NAME)." ) resolution_subtype: str | None = Field( None, description="Resolution subtype (e.g. RESOLUTION_SUBTYPE_NAME)." ) YouTubeAssetID = Annotated[str, StringConstraints(pattern=r".{5,}")] class YTListRequest(BaseModel): """YouTube Content ID List request.""" asset_ids: list[YouTubeAssetID] = Field( ..., description="List of YouTube Content ID Assets.", min_items=1 ) class ExportSheet(BaseModel): """Export rows to Excel byte stream.""" rows: list[dict] = Field(..., description="List of rows.") columns: list[list[str]] | None = OptionalField( description="List of columns in the desired order. Each column is a tuple of " "column key (relating it to the row keys) and column label (as it will " "appear in the Excel file)." ) name: str | None = OptionalField(description="Sheet name.") class Export(BaseModel): """Export sheets to Excel byte stream.""" sheets: list[ExportSheet] = Field( ..., description="List of sheets to export. Each sheet is a list of rows which will " "be exported to a separate sheet in the Excel file.", ) class AuditReportRows(BaseModel): """Audit report rows model.""" audit_group_id: AuditGroupID = Field(..., description="Audit group ID.") class AuditReport(BaseModel): """Audit report data model.""" audit_group_id: AuditGroupID = Field(..., description="Audit group ID.") label_name: str = Field(..., description="Label name.") artist_image: str | None = OptionalField( description="Artist image URL, URL data or None. If a URL is provided, the " "image will be fetched and included in the report. If a data URL (i.e. " "base64-encoded image) is provided, the image will be included directly in " "the report. The provided data URL must be in the format " "'data:image/...;base64, ...'." ) as_pdf: bool = Field(True, description="Whether to output the report as PDF.") class AuditMetaSet(BaseModel): """Set audit metadata model.""" audit_id: AuditID = Field(..., description="Audit ID.") keys: str | None = OptionalField( description="Metadata keys to filter. If None, " "all keys." ) value: Any = Field( ..., description="Metadata value. Anything other than None will be converted to string.", ) @dataclass class NewAuditFlag: """New audit flag, for DB insertion. Setting the `resolution` field will auto-resolve the flag. Args: row_idx: Row index. isrc: ISRC. upc: UPC. asset_id: Asset ID. video_id: Video ID. text: Flag text (e.g. "MUST_CLAIM") details: Optional flag details for context in the UI. resolution: Resolution (e.g. "IGNORE") resolution_subtype: Resolution subtype (e.g. "ASSET_OK") """ text: FlagText row_idx: int isrc: ISRC | None = None upc: UPC | None = None asset_id: AssetID | None = None video_id: VideoID | None = None details: Any = None # It is possible to add auto-resolved flags, if at least `resolution` is not None. resolution: str | None = None resolution_subtype: str | None = None def __post_init__(self): # Convert np.nan to None for all fields for field_name, field_value in self.__dict__.items(): if isinstance(field_value, float) and np.isnan(field_value): setattr(self, field_name, None) try: self.text = self.text.strip() except AttributeError: raise ValueError("Flag text cannot be empty.") from None if self.isrc is None and self.upc is None and self.video_id is None: raise ValueError( "At least one of ISRC, UPC or Video ID must be provided, or several." ) if isinstance(self.upc, int): self.upc = str(self.upc) if not isinstance(self.details, str) and self.details is not None: if isinstance(self.details, (list, tuple, set)): self.details = ", ".join(sorted(list(self.details))) else: self.details = str(self.details) if self.details and isinstance(self.details, str): self.details.strip() if self.resolution and not isinstance(self.resolution, FlagResolutions): raise ValueError("Invalid resolution, must be a FlagResolution.") if self.resolution_subtype and not isinstance( self.resolution_subtype, FlagResolutions ): raise ValueError("Invalid resolution subtype, must be a FlagResolution.") if self.resolution_subtype is not None and self.resolution is None: raise ValueError("Resolution subtype cannot be set without resolution.")