diff --git a/backend/apps/library/services.py b/backend/apps/library/services.py index 635ff37..485ed0a 100644 --- a/backend/apps/library/services.py +++ b/backend/apps/library/services.py @@ -231,6 +231,7 @@ def trim_e621_post(post): "size": file_data.get("size"), "width": file_data.get("width"), "height": file_data.get("height"), + "url": file_data.get("url"), }, "uploader_name": post.get("uploader_name"), } @@ -244,12 +245,38 @@ def sanitize_iqdb_results(results): continue score = result.get("score") post_id = result.get("post_id") + tags_preview = result.get("tags_preview") entry = { "post_id": post_id if isinstance(post_id, int) else None, "score": float(score) if isinstance(score, (int, float)) else None, "preview_url": str(result.get("preview_url") or "")[:500] or None, "rating": str(result.get("rating") or "")[:1] or None, "md5": str(result.get("md5") or "")[:32] or None, + "score_total": ( + result.get("score_total") + if isinstance(result.get("score_total"), int) + else None + ), + "fav_count": ( + result.get("fav_count") + if isinstance(result.get("fav_count"), int) + else None + ), + "width": ( + result.get("width") + if isinstance(result.get("width"), int) + else None + ), + "height": ( + result.get("height") + if isinstance(result.get("height"), int) + else None + ), + "tags_preview": ( + [str(tag)[:100] for tag in tags_preview][:8] + if isinstance(tags_preview, list) + else [] + ), } if entry["post_id"] or entry["preview_url"]: cleaned.append(entry) diff --git a/backend/apps/library/uploads.py b/backend/apps/library/uploads.py index cd53eef..6480e25 100644 --- a/backend/apps/library/uploads.py +++ b/backend/apps/library/uploads.py @@ -12,6 +12,7 @@ Flow: import logging import shutil from pathlib import Path +from urllib.parse import urlparse from django.conf import settings from django.contrib.auth import get_user_model @@ -31,20 +32,43 @@ from .serializers import TempUploadSerializer logger = logging.getLogger(__name__) -def complete_temp_upload(temp): - """Move a staged file into the watched folder and index it.""" +def complete_temp_upload(temp, download_url=None): + """Index the upload into the library. + + With a ``download_url`` the e621 original is fetched and the staged copy + is dropped; otherwise the staged file itself moves into the library. + """ folder = Path(settings.WATCHED_FOLDER) - if not temp.file: - raise ValueError("staged file is missing") - destination = services.unique_destination( - folder, temp.original_filename or "upload" - ) - with temp.file.open("rb") as source, open(destination, "wb") as target: - shutil.copyfileobj(source, target) + if download_url: + file_data = {} + if isinstance(temp.e621_data, dict): + file_data = temp.e621_data.get("file") or {} + extension = Path(urlparse(download_url).path).suffix + if not extension: + extension = Path(temp.original_filename).suffix + md5 = file_data.get("md5") + name = f"{md5}{extension}" if md5 else (temp.original_filename or "download") + destination = services.unique_destination(folder, name) + try: + services.download_file(download_url, destination) + except Exception: + destination.unlink(missing_ok=True) + raise + item, _, _ = services.index_file(destination, folder) + if temp.file: + temp.file.delete(save=False) + else: + if not temp.file: + raise ValueError("staged file is missing") + destination = services.unique_destination( + folder, temp.original_filename or "upload" + ) + with temp.file.open("rb") as source, open(destination, "wb") as target: + shutil.copyfileobj(source, target) + item, _, _ = services.index_file(destination, folder) + temp.file.delete(save=False) - item, _, _ = services.index_file(destination, folder) - temp.file.delete(save=False) temp.library_item = item temp.status = TempUpload.STATUS_COMPLETED @@ -184,6 +208,7 @@ class TempUploadViewSet( ) mode = str(request.data.get("mode") or "").strip() + download_url = None if mode == "link": post = request.data.get("post") post_id = request.data.get("post_id") @@ -213,6 +238,17 @@ class TempUploadViewSet( notes = request.data.get("notes") if notes is not None: temp.custom_notes = str(notes) + + # Prefer the e621 original; fall back to the staged file when the + # post has no URL or its file is byte-identical to the upload. + file_data = (trimmed or {}).get("file") or {} + candidate_url = str( + request.data.get("file_url") or file_data.get("url") or "" + ).strip() + post_md5 = str(file_data.get("md5") or "").strip().lower() + if candidate_url and post_md5 and post_md5 == temp.md5.lower(): + candidate_url = "" + download_url = candidate_url or None elif mode == "custom": rating = str(request.data.get("rating") or "") temp.custom_rating = rating if rating in {"s", "q", "e"} else "" @@ -227,7 +263,7 @@ class TempUploadViewSet( temp.save() try: - complete_temp_upload(temp) + complete_temp_upload(temp, download_url=download_url) except Exception as exc: # noqa: BLE001 - report completion failures logger.exception("Could not complete staged upload %s", temp.id) temp.status = TempUpload.STATUS_ERROR diff --git a/frontend/src/features/upload/UploadPage.tsx b/frontend/src/features/upload/UploadPage.tsx index d483ea8..93711a7 100644 --- a/frontend/src/features/upload/UploadPage.tsx +++ b/frontend/src/features/upload/UploadPage.tsx @@ -17,6 +17,7 @@ import { fetchPost, fetchPostsByMd5, iqdbSearch, + RATING_LABELS, } from "@/lib/e621"; import { formatBytes } from "@/lib/format"; import type { @@ -45,6 +46,12 @@ const RESOLUTION_LABELS: Record = { custom: "custom metadata", }; +const ratingPill: Record = { + s: "bg-ctp-green text-ctp-crust", + q: "bg-ctp-peach text-ctp-crust", + e: "bg-ctp-red text-ctp-crust", +}; + const ratingChoices: { value: Rating; label: string; active: string }[] = [ { value: "", @@ -205,6 +212,7 @@ function MetadataModal({ const credentials = useE621((state) => state.credentials); const [tab, setTab] = useState<"link" | "custom">("link"); const [postId, setPostId] = useState(""); + const [selected, setSelected] = useState(null); const [rating, setRating] = useState(temp.custom_rating); const [tags, setTags] = useState(temp.custom_tags.join(", ")); const [notes, setNotes] = useState(temp.custom_notes); @@ -354,17 +362,39 @@ function MetadataModal({ key={`${candidate.post_id}-${candidate.md5}`} type="button" disabled={busy || !candidate.post_id} - onClick={() => void linkCandidate(candidate)} - className="overflow-hidden rounded-md border border-ctp-surface0 bg-ctp-mantle text-left transition hover:border-ctp-surface2 disabled:opacity-50" + onClick={() => setSelected(candidate)} + className={cn( + "overflow-hidden rounded-md border bg-ctp-mantle text-left transition disabled:opacity-50", + selected?.post_id === candidate.post_id + ? "border-ctp-mauve ring-2 ring-ctp-mauve/40" + : "border-ctp-surface0 hover:border-ctp-surface2", + )} > - {candidate.preview_url ? ( - {`Post - ) : null} + + {candidate.preview_url ? ( + {`Post + ) : ( + + No preview + + )} + {candidate.rating ? ( + + {candidate.rating} + + ) : null} + #{candidate.post_id} {candidate.score !== null @@ -374,6 +404,62 @@ function MetadataModal({ ))} + + {selected ? ( +
+

+ Post #{selected.post_id} +

+

+ {selected.rating + ? (RATING_LABELS[selected.rating] ?? selected.rating) + : "Unknown rating"} + {selected.score_total !== null && + selected.score_total !== undefined + ? ` · ▲ ${selected.score_total}` + : ""} + {selected.fav_count !== null && + selected.fav_count !== undefined + ? ` · ${selected.fav_count} favs` + : ""} + {selected.width && selected.height + ? ` · ${selected.width}×${selected.height}` + : ""} +

+ {selected.tags_preview && + selected.tags_preview.length > 0 ? ( +
+ {selected.tags_preview.map((tag) => ( + + {tag} + + ))} +
+ ) : null} +
+ + + View post + +
+

+ Linking downloads the post's file into the library and + drops this staged upload. +

+
+ ) : null} ) : (

diff --git a/frontend/src/lib/e621.ts b/frontend/src/lib/e621.ts index 755a6d1..a0eef1c 100644 --- a/frontend/src/lib/e621.ts +++ b/frontend/src/lib/e621.ts @@ -322,6 +322,26 @@ export interface E621IqdbResult { preview_url: string | null; rating: string | null; md5: string | null; + score_total: number | null; + fav_count: number | null; + width: number | null; + height: number | null; + tags_preview: string[]; +} + +function flattenTagPreview(tags: unknown, limit = 8): string[] { + if (!tags || typeof tags !== "object") return []; + const out: string[] = []; + for (const value of Object.values(tags as Record)) { + if (!Array.isArray(value)) continue; + for (const tag of value) { + if (typeof tag === "string" && !out.includes(tag)) { + out.push(tag); + if (out.length >= limit) return out; + } + } + } + return out; } function normalizeIqdbEntry(entry: unknown): E621IqdbResult | null { @@ -344,12 +364,18 @@ function normalizeIqdbEntry(entry: unknown): E621IqdbResult | null { : typeof post.id === "number" ? post.id : null; + const scoreTotal = post.score; return { post_id: postId, score: typeof record.score === "number" ? record.score : null, preview_url: typeof preview.url === "string" ? preview.url : null, rating: typeof post.rating === "string" ? post.rating : null, md5: typeof file.md5 === "string" ? file.md5 : null, + score_total: typeof scoreTotal === "number" ? scoreTotal : null, + fav_count: typeof post.fav_count === "number" ? post.fav_count : null, + width: typeof file.width === "number" ? file.width : null, + height: typeof file.height === "number" ? file.height : null, + tags_preview: flattenTagPreview(post.tags), }; } @@ -371,13 +397,44 @@ export async function iqdbSearch( : Array.isArray((data as { matches?: unknown[] })?.matches) ? (data as { matches: unknown[] }).matches : []; - return list + const candidates = list .map(normalizeIqdbEntry) .filter( (entry): entry is E621IqdbResult => entry !== null && (entry.post_id !== null || entry.preview_url !== null), ) .slice(0, 10); + + // The IQDB payload carries no preview/file data; enrich with one batched + // post lookup so candidates can show what they actually are. + const ids = candidates + .map((candidate) => candidate.post_id) + .filter((id): id is number => typeof id === "number"); + if (ids.length > 0) { + try { + const posts = await fetchPostsByIds(credentials, ids, signal); + const byId = new Map(posts.map((post) => [post.id, post])); + return candidates.map((candidate) => { + const post = + candidate.post_id !== null ? byId.get(candidate.post_id) : undefined; + if (!post) return candidate; + return { + ...candidate, + preview_url: post.preview.url ?? candidate.preview_url, + rating: post.rating ?? candidate.rating, + md5: post.file.md5 ?? candidate.md5, + score_total: post.score.total, + fav_count: post.fav_count, + width: post.file.width, + height: post.file.height, + tags_preview: flattenTagPreview(post.tags), + }; + }); + } catch { + // Keep the base candidates when the enrichment request fails. + } + } + return candidates; } export function postIsBlacklisted( diff --git a/frontend/src/lib/types.ts b/frontend/src/lib/types.ts index 85639be..30813ec 100644 --- a/frontend/src/lib/types.ts +++ b/frontend/src/lib/types.ts @@ -113,6 +113,11 @@ export interface E621IqdbCandidate { preview_url: string | null; rating: string | null; md5: string | null; + score_total?: number | null; + fav_count?: number | null; + width?: number | null; + height?: number | null; + tags_preview?: string[]; } export interface TempUpload {