feat(collections): add collection routes and the Images collection filter

This commit is contained in:
Vito Sansevero
2026-10-27 18:02:29 -07:00
parent b543b7db89
commit cb874d7826
11 changed files with 1182 additions and 80 deletions
+3
View File
@@ -48,6 +48,7 @@ from .review import ReviewRoutesMixin
from .review_recording import ReviewRecordingMixin
from .review_events import ReviewEventsMixin
from .saved_views import SavedViewRoutesMixin
from .collections import CollectionRoutesMixin
from .history import HistoryRoutesMixin
from .images import ImageRoutesMixin
from .gallery_listing import WALK_STOPPED, RootListing, WalkOutcome
@@ -314,6 +315,7 @@ class PromptManagerAPI(
ReviewRecordingMixin,
ReviewEventsMixin,
SavedViewRoutesMixin,
CollectionRoutesMixin,
HistoryRoutesMixin,
PromptRoutesMixin,
ImageRoutesMixin,
@@ -684,6 +686,7 @@ class PromptManagerAPI(
self._register_review_routes(routes)
self._register_review_event_routes(routes)
self._register_saved_view_routes(routes)
self._register_collection_routes(routes)
self._register_history_routes(routes)
self._register_prompt_routes(routes)
self._register_image_routes(routes)
+8 -36
View File
@@ -27,7 +27,12 @@ from .annotation_suggest import (
validate_options,
)
from .autotag_routes import model_busy_response
from .images import _forbidden, _has_traversal, _unavailable_link, _verified_title
# The target check is shared with the collection routes (#331); both names
# stay importable from here.
from .media_target import MAX_TARGET_PATH # noqa: F401
from .media_target import media_target as _target
from .images import _unavailable_link, _verified_title
from .prompts import bad_request
from .. import model_gate
@@ -40,7 +45,7 @@ try:
)
from ...database.execution_history import subscribe_history_restore
from ...database.library_commands import LibraryCommandError
from ...utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore
from ...utils.media_capture import CaptureError, CaptureStore
except ImportError:
from database.annotations import (
MAX_IMAGE_TAGS,
@@ -50,15 +55,13 @@ except ImportError:
)
from database.execution_history import subscribe_history_restore
from database.library_commands import LibraryCommandError
from utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore
from utils.media_capture import CaptureError, CaptureStore
# Request bodies of the annotation routes are small (a target plus a tag set).
MAX_BODY_BYTES = 64 * 1024
MAX_TARGET_PATH = 1024
CAPTURE_URL = "/prompt_manager/image-annotations/capture/"
_PREVIEW_CHUNK = 64 * 1024
_TARGET_FIELDS = frozenset({"root_index", "root_scope", "path"})
_OBSERVE_FIELDS = frozenset({"target"})
_SAVE_FIELDS = frozenset(
{
@@ -188,37 +191,6 @@ def _capture_error(exc: CaptureError) -> web.Response:
return _coded(exc.code, _MESSAGES.get(exc.code) or str(exc))
def _target(data: object):
"""``((index, scope, relative), None)`` for a target, else ``(None, error)``.
Only the listing identity is accepted: an int root index, the root's
64-hex scope and a bounded root-relative POSIX media path, never an
absolute, drive-qualified, backslash or ``..`` path. Nothing here touches
the filesystem or the database.
"""
if not isinstance(data, dict) or set(data) != _TARGET_FIELDS:
return None, bad_request("target must name root_index, root_scope and path")
index, scope, relative = data["root_index"], data["root_scope"], data["path"]
if type(index) is not int or index < 0:
return None, bad_request("Invalid gallery root")
if not isinstance(scope, str) or _HEX64.fullmatch(scope) is None:
return None, bad_request("Invalid gallery root scope")
if not isinstance(relative, str) or not relative or len(relative) > MAX_TARGET_PATH:
return None, bad_request("path must be a relative media path")
parts = relative.split("/")
if (
"\0" in relative
or "\\" in relative
or relative.startswith("/")
or _has_traversal(relative)
or any(part in ("", ".", "..") or ":" in part for part in parts)
):
return None, _forbidden()
if os.path.splitext(parts[-1])[1].lower() not in CAPTURE_EXTENSIONS:
return None, _forbidden("Only media files can be tagged")
return (index, scope, relative), None
def _expires_at(seconds: float | None) -> str:
moment = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
seconds=max(0.0, seconds or 0.0)
+325
View File
@@ -0,0 +1,325 @@
"""Collections (#331): named, hand-picked groups of images.
Routes under ``/prompt_manager/collections``. A file is named by its listing
identity ``{root_index, root_scope, path}`` (``media_target.py``) and stored
under the favorites' keys (``_favorite_root_key``, ``_favorite_relative_key``),
which never leave the server. A stale scope (the sources were reordered) is a
409 ``source-changed`` before anything is written. Responses are
``{"success": bool, ...}``; an unexpected failure is a logged 500 with a
generic message, never the exception text.
"""
import hmac
import re
from functools import wraps
from typing import Callable
from aiohttp import web
from .images import (
_favorite_relative_key,
_favorite_root_key,
_forbidden,
_json_body_or_error,
collection_not_found,
)
from .media_target import media_target
from .prompts import bad_request
try:
from ...database.image_collections import (
CollectionFilesMissing,
CollectionInvalid,
CollectionMissing,
CollectionNameTaken,
normalize_collection_name,
)
from ...database.image_favorites import MAX_KEY_LENGTH
except ImportError:
from database.image_collections import (
CollectionFilesMissing,
CollectionInvalid,
CollectionMissing,
CollectionNameTaken,
normalize_collection_name,
)
from database.image_favorites import MAX_KEY_LENGTH
BASE = "/prompt_manager/collections"
# The most files one add or remove names (the bulk tagging cap, MC12).
MAX_ITEMS = 500
MAX_ID = 999_999_999
_PATH_ID = re.compile(r"[0-9]{1,9}")
_ROOT_INDEX = re.compile(r"[0-9]{1,9}")
_ADD_TARGETS = (frozenset({"collection_id", "items"}), frozenset({"name", "items"}))
_REMOVE_FIELDS = frozenset({"collection_id", "items"})
_RENAME_FIELDS = frozenset({"name"})
_MEDIA_MESSAGE = "Only media files can be added to a collection"
READ_FAILED = "Collections could not be read."
CHANGE_FAILED = "The collection could not be changed."
SOURCES_CHANGED = "The sources changed. Refresh the images and try again."
FILES_NOT_FOUND = (
"None of the selected images could be found. Refresh the images and try again."
)
NAME_TAKEN = "Another collection already has this name."
class _Refused(Exception):
"""A request refused before any write; carries its response."""
def __init__(self, response: web.Response):
super().__init__(response.status)
self.response = response
def _coded(status: int, code: str, message: str, **extra) -> web.Response:
return web.json_response(
{"success": False, "code": code, "error": message, **extra},
status=status,
headers={"Cache-Control": "no-store"},
)
def _collection_response(failure: str) -> Callable:
"""Map the database's refusals to their status; anything else to a 500."""
def decorate(handler: Callable) -> Callable:
@wraps(handler)
async def wrapped(self, request: web.Request) -> web.Response:
try:
return await handler(self, request)
except _Refused as refused:
return refused.response
except CollectionInvalid as exc:
return bad_request(str(exc))
except CollectionMissing:
return collection_not_found()
except CollectionFilesMissing:
return _coded(404, "files-not-found", FILES_NOT_FOUND)
except CollectionNameTaken as exc:
return _coded(
409, "name-taken", NAME_TAKEN, collection_id=exc.collection_id
)
except Exception:
self.logger.exception("Collection request failed")
return web.json_response(
{"success": False, "error": failure}, status=500
)
return wrapped
return decorate
def _path_id(request: web.Request) -> int:
raw = request.match_info["collection_id"]
if not raw.isascii() or _PATH_ID.fullmatch(raw) is None or int(raw) < 1:
raise _Refused(bad_request("Invalid collection id"))
return int(raw)
def _body_id(value: object) -> int:
if type(value) is not int or not 1 <= value <= MAX_ID:
raise _Refused(bad_request("Invalid collection id"))
return value
async def _body(request: web.Request, allowed: tuple) -> dict:
data, error = await _json_body_or_error(request)
if error is not None:
raise _Refused(error)
if frozenset(data) not in allowed:
names = " or ".join(", ".join(sorted(fields)) for fields in allowed)
raise _Refused(bad_request(f"Request must contain exactly {names}"))
return data
class CollectionRoutesMixin:
"""Mixin providing the collection routes (#331)."""
def _register_collection_routes(self, routes: web.RouteTableDef) -> None:
# The fixed paths go first: "/items" must never match "/{collection_id}".
routes.get(BASE)(self.list_collections)
routes.post(BASE + "/items")(self.add_collection_items)
routes.delete(BASE + "/items")(self.remove_collection_items)
routes.get(BASE + "/membership")(self.collection_membership)
routes.patch(BASE + "/{collection_id}")(self.rename_collection)
routes.delete(BASE + "/{collection_id}")(self.delete_collection)
# -- reads -------------------------------------------------------------
@_collection_response(READ_FAILED)
async def list_collections(self, request: web.Request) -> web.Response:
"""Every collection with its listed ``count`` and its ``missing`` rows.
``count`` is the members listed under every configured source (one
primary-key range read per source, matched against its cached
listing); ``missing`` the stored rows beyond it. When a scan or a read
fails both are null for every collection and the names still return.
"""
collections = await self._run_in_executor(self.db.list_collections)
counts = None
if collections:
try:
counts = await self._collection_counts()
except Exception:
self.logger.warning(
"Collection counts could not be read", exc_info=True
)
rows = []
for collection in collections:
listed, stored = (None, None) if counts is None else counts
count = None if listed is None else listed.get(collection["id"], 0)
missing = None if stored is None else stored.get(collection["id"], 0)
rows.append(
{
**collection,
"count": count,
"missing": None if count is None else missing - count,
}
)
return web.json_response({"success": True, "collections": rows})
async def _collection_counts(self) -> tuple[dict, dict]:
"""``({id: listed members}, {id: stored rows})`` over every source."""
listed: dict[int, int] = {}
seen = set()
for root in self._get_all_output_dirs():
root_key = _favorite_root_key(root)
if root_key in seen: # one folder configured twice counts once
continue
seen.add(root_key)
files = await self._get_gallery_files(root)
listing = await self._run_in_executor(self._favorite_listing, root, files)
members = await self._run_in_executor(
self.db.root_collection_members, root_key
)
for collection_id, keys in members.items():
found = sum(1 for key in keys if key in listing)
listed[collection_id] = listed.get(collection_id, 0) + found
stored = await self._run_in_executor(self.db.collection_counts)
return listed, stored
@_collection_response(READ_FAILED)
async def collection_membership(self, request: web.Request) -> web.Response:
"""GET ``?root_index&root_scope&path``: the ids holding that file."""
query = request.query
if set(query) != {"root_index", "root_scope", "path"}:
return bad_request("Query must name root_index, root_scope and path")
raw_index = query["root_index"]
if not raw_index.isascii() or _ROOT_INDEX.fullmatch(raw_index) is None:
return bad_request("Invalid gallery root")
roots = list(self._get_all_output_dirs())
root_key, key = self._collection_key(
{
"root_index": int(raw_index),
"root_scope": query["root_scope"],
"path": query["path"],
},
roots,
)
ids = await self._run_in_executor(self.db.file_collections, root_key, key)
return web.json_response({"success": True, "collection_ids": ids})
# -- writes ------------------------------------------------------------
@_collection_response(CHANGE_FAILED)
async def add_collection_items(self, request: web.Request) -> web.Response:
"""POST ``{collection_id | name, items}``: add the listed files.
Only files the source's listing holds are added; the rest count in
``not_found``. A new name with nothing to add creates nothing (404
``files-not-found``).
"""
data = await _body(request, _ADD_TARGETS)
target = {}
if "name" in data:
target["name"] = normalize_collection_name(data["name"])
else:
target["collection_id"] = _body_id(data["collection_id"])
roots = list(self._get_all_output_dirs())
keys = self._collection_items(data["items"], roots)
found, not_found = [], 0
listings = {}
for index, root_key, key in keys:
if index not in listings:
root = roots[index]
files = await self._get_gallery_files(root)
listings[index] = await self._run_in_executor(
self._favorite_listing, root, files
)
if key in listings[index]:
found.append((root_key, key))
else:
not_found += 1
result = await self._run_in_executor(
lambda: self.db.add_to_collection(items=found, **target)
)
return web.json_response({"success": True, **result, "not_found": not_found})
@_collection_response(CHANGE_FAILED)
async def remove_collection_items(self, request: web.Request) -> web.Response:
"""DELETE ``{collection_id, items}``: listed or not, every named file."""
data = await _body(request, (_REMOVE_FIELDS,))
collection_id = _body_id(data["collection_id"])
roots = list(self._get_all_output_dirs())
keys = self._collection_items(data["items"], roots)
result = await self._run_in_executor(
self.db.remove_from_collection,
collection_id,
[(root_key, key) for _index, root_key, key in keys],
)
return web.json_response({"success": True, **result})
@_collection_response(CHANGE_FAILED)
async def rename_collection(self, request: web.Request) -> web.Response:
"""PATCH ``{name}``: a case variant of its own name is allowed."""
collection_id = _path_id(request)
data = await _body(request, (_RENAME_FIELDS,))
name = normalize_collection_name(data["name"])
collection = await self._run_in_executor(
self.db.rename_collection, collection_id, name
)
return web.json_response({"success": True, "collection": collection})
@_collection_response(CHANGE_FAILED)
async def delete_collection(self, request: web.Request) -> web.Response:
"""DELETE: the collection and its members; files are never touched."""
collection_id = _path_id(request)
await self._run_in_executor(self.db.delete_collection, collection_id)
return web.json_response({"success": True, "deleted": True})
# -- request steps -----------------------------------------------------
def _collection_items(self, items: object, roots: list) -> list[tuple]:
"""``(index, root key, relative key)`` per distinct file, in order.
Every item is checked before anything is written: 400 for a malformed
shape, 403 for an unsafe path, 409 for a stale scope. Two spellings
that fold to one key (Windows) count once.
"""
if not isinstance(items, list) or not 1 <= len(items) <= MAX_ITEMS:
raise _Refused(bad_request(f"items must list 1 to {MAX_ITEMS} images"))
keys = {}
for item in items:
index = item.get("root_index") if isinstance(item, dict) else None
root_key, key = self._collection_key(item, roots)
keys.setdefault((root_key, key), index)
return [(index, root_key, key) for (root_key, key), index in keys.items()]
def _collection_key(self, item: object, roots: list) -> tuple[str, str]:
"""The stored keys of one target, or raise its refusal."""
target, error = media_target(item, media_message=_MEDIA_MESSAGE)
if error is not None:
raise _Refused(error)
index, scope, relative = target
if index >= len(roots) or not hmac.compare_digest(
scope, self._gallery_root_scope(roots[index], index)
):
raise _Refused(_coded(409, "source-changed", SOURCES_CHANGED))
root_key = _favorite_root_key(roots[index])
key = _favorite_relative_key(relative)
if max(len(root_key), len(key)) > MAX_KEY_LENGTH:
# A stored key past the schema bound would fail the next restore.
raise _Refused(_forbidden("Path is too long to add to a collection"))
return root_key, key
+4
View File
@@ -58,7 +58,9 @@ from .annotations import (
from .autotag_routes import BATCH_CANCELLED, MODEL_BUSY_ERROR, model_busy_response
from .images import (
MAX_LISTING_SELECTION,
CollectionMissing,
ListingSourceChanged,
collection_not_found,
_current_fingerprint,
_listing_posix,
_public_path,
@@ -685,6 +687,8 @@ class ImageTagJobsMixin:
selection = await self.resolve_listing_filter(spec.filter, cap)
except ListingSourceChanged:
return None, _coded("source-changed")
except CollectionMissing:
return None, collection_not_found()
except ValueError as exc:
return None, bad_request(str(exc))
except Exception:
+110 -44
View File
@@ -18,7 +18,7 @@ import urllib.parse
from bisect import bisect_left, bisect_right, insort
from collections.abc import Mapping
from pathlib import Path
from typing import Callable, NamedTuple
from typing import Callable, Iterable, NamedTuple
from aiohttp import web
from PIL import Image, UnidentifiedImageError
@@ -58,6 +58,7 @@ from .inspection_service import (
)
try:
from ...database.image_collections import CollectionMissing
from ...database.image_favorites import MAX_KEY_LENGTH
from ...database.library import display_title
from ...database.models import normalize_image_path
@@ -78,6 +79,7 @@ try:
reveal_in_file_manager,
)
except ImportError:
from database.image_collections import CollectionMissing
from database.image_favorites import MAX_KEY_LENGTH
from database.library import display_title
from database.models import normalize_image_path
@@ -104,6 +106,21 @@ IMAGE_EXTENSIONS = _IMAGE_EXTENSIONS
_CARD_RECORDS = ("prompt", "workflow")
_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
_LISTING_FAILED = "The image list could not be filtered. Check server logs for details."
_COLLECTION_ID = re.compile(r"[1-9][0-9]{0,8}")
COLLECTION_NOT_FOUND = "That collection no longer exists."
def collection_not_found() -> web.Response:
"""The 404 a write or a listing answers for an unknown collection (#331)."""
return web.json_response(
{
"success": False,
"code": "collection-not-found",
"error": COLLECTION_NOT_FOUND,
},
status=404,
headers={"Cache-Control": "no-store"},
)
def _card_header(path: Path) -> tuple:
@@ -1056,8 +1073,9 @@ def _workspace_filters(query: Mapping) -> dict:
the linked positive's text, title and tag names. ``category`` is exact,
``tags`` is a comma list that must all be present, ``sort=rating`` ranks by
the linked positive, and ``favorites=1`` keeps only favorites (the
Favorites page, #249 D51). Raises ValueError for any unknown or oversized
value.
Favorites page, #249 D51). ``collection=<id>`` (1 to 9 digits) keeps only
that collection's members (#331); it wins over ``favorites`` (a page never
sets both). Raises ValueError for any unknown or oversized value.
"""
text = query.get("q", "").strip()
folded = text.casefold()
@@ -1071,8 +1089,10 @@ def _workspace_filters(query: Mapping) -> dict:
{tag.strip() for tag in query.get("tags", "").split(",") if tag.strip()}
)
favorites = query.get("favorites", "0")
collection = query.get("collection", "")
if (
favorites not in ("0", "1")
or (collection != "" and _COLLECTION_ID.fullmatch(collection) is None)
or len(text) > 500
or len(folded) > 500
or media not in ("all", "image", "video")
@@ -1095,7 +1115,8 @@ def _workspace_filters(query: Mapping) -> dict:
"path": exact.casefold(),
"category": category,
"tags": tags,
"favorites": favorites == "1",
"favorites": favorites == "1" and not collection,
"collection": int(collection) if collection else None,
}
@@ -1113,6 +1134,8 @@ LISTING_FILTER_KEYS = frozenset(
"media",
"sort",
"favorites",
# A collection page's filter (#331): bulk actions act on its members.
"collection",
}
)
MAX_LISTING_FILTER_VALUE = 2048
@@ -1169,7 +1192,7 @@ def _admitted_entries(
entries: list,
filters: dict,
links: dict | None,
favorites: frozenset | None = None,
members: frozenset | None = None,
image_tags: dict | None = None,
):
"""Yield ``(position, entry, link match)`` for each entry the filters admit.
@@ -1177,11 +1200,11 @@ def _admitted_entries(
*links* maps normalized paths to (hit, rating, positive id) for every
unambiguously linked image (see PromptDatabase.get_image_link_matches); a
category or tag filter admits only those, and ``q`` admits a filename match
or a linked-positive hit. *favorites*, when given, admits only those listed
paths (``favorites=1``). *image_tags* maps normalized paths to their
_ImageTagMatch: a tag filter also admits a file whose current image tags
(with its linked positive's tags) carry every tag, and ``q`` also admits an
image tag name hit.
or a linked-positive hit. *members*, when given, admits only those listed
paths (the favorites of ``favorites=1`` or a ``collection``'s members).
*image_tags* maps normalized paths to their _ImageTagMatch: a tag filter
also admits a file whose current image tags (with its linked positive's
tags) carry every tag, and ``q`` also admits an image tag name hit.
"""
query, media, folder = filters["query"], filters["media"], filters["folder"]
members_only = filters["category"] is not None or bool(filters["tags"])
@@ -1197,7 +1220,7 @@ def _admitted_entries(
continue
if media != "all" and (media == "video") != entry.video:
continue
if favorites is not None and entry.path not in favorites:
if members is not None and entry.path not in members:
continue
match = None if links is None else links.get(entry.normalized)
image = None if image_tags is None else image_tags.get(entry.normalized)
@@ -1219,7 +1242,7 @@ def _filtered_page(
size: int,
filters: dict,
links: dict | None,
favorites: frozenset | None = None,
members: frozenset | None = None,
image_tags: dict | None = None,
) -> tuple[int, list, int | None]:
"""Count the *entries* the filters admit and take *size* from *start* on.
@@ -1230,7 +1253,7 @@ def _filtered_page(
"""
total, page, positives = 0, [], set()
for position, entry, match in _admitted_entries(
entries, filters, links, favorites, image_tags
entries, filters, links, members, image_tags
):
total += 1
if match is not None:
@@ -1934,6 +1957,7 @@ class ImageRoutesMixin:
filters["category"],
filters["tags"],
filters["favorites"],
filters["collection"],
]
).encode()
).hexdigest()
@@ -1944,17 +1968,19 @@ class ImageRoutesMixin:
except ValueError as exc:
return bad_request(str(exc))
listings = [
(index, root, await self._get_gallery_files(root))
for index, root in enumerate(output_dirs)
if root_index is None or index == root_index
]
try:
if filters["collection"] is not None and not await self._run_in_executor(
self.db.collection_exists, filters["collection"]
):
raise CollectionMissing("No such collection")
listings = await self._workspace_listings(output_dirs, root_index)
page, total, has_more, next_cursor, positive_count = (
await self._run_in_executor(
self._select_workspace_page, listings, filters, after, limit, scope
)
)
except CollectionMissing:
return collection_not_found()
except Exception:
# A filter that needs the prompt links or the image tags cannot be
# served without them. Never echo the exception: it can carry SQL
@@ -2215,9 +2241,11 @@ class ImageRoutesMixin:
) -> tuple:
"""The ordered entries and what admits them (blocking; run in executor).
Returns ``(entries, key, links, image_tags, favorites)``: the listing in
the filters' order, its keyset key, and the link, image-tag and favorite
maps _admitted_entries reads. One SELECT (get_image_link_matches) reads
Returns ``(entries, key, links, image_tags, members)``: the listing in
the filters' order, its keyset key, and the link, image-tag and member
maps _admitted_entries reads (members: a collection's, else the
favorites, else None). A collection deleted since the route checked it
raises CollectionMissing. One SELECT (get_image_link_matches) reads
every linked positive when *read_links*. A failed read raises when the
filters need it (category, tags, the rating sort) or *strict*; else the
listing stays usable without links (a search then matches filenames
@@ -2290,10 +2318,13 @@ class ImageRoutesMixin:
def key(entry):
return _workspace_key(entry, order)
favorites = (
self._workspace_favorite_paths(listings) if filters["favorites"] else None
return (
entries,
key,
links,
image_tags,
self._workspace_members(listings, filters),
)
return entries, key, links, image_tags, favorites
def _select_workspace_page(
self,
@@ -2322,7 +2353,7 @@ class ImageRoutesMixin:
or bool(filters["tags"])
or order == "rating"
)
entries, key, links, image_tags, favorites = self._workspace_admission(
entries, key, links, image_tags, members = self._workspace_admission(
listings, filters, read_links=needed or after is None
)
start = 0
@@ -2333,7 +2364,7 @@ class ImageRoutesMixin:
else:
start = bisect_right(entries, after, key=key)
total, page, positive_count = _filtered_page(
entries, start, limit + 1, filters, links, favorites, image_tags
entries, start, limit + 1, filters, links, members, image_tags
)
if not needed and after is not None:
positive_count = _COUNT_OMITTED
@@ -2355,12 +2386,12 @@ class ImageRoutesMixin:
The links are read strictly: the positives must be the listing's own.
"""
entries, _key, links, image_tags, favorites = self._workspace_admission(
entries, _key, links, image_tags, members = self._workspace_admission(
listings, filters, read_links=True, strict=True
)
targets, positives, unlinked, scopes = [], {}, 0, {}
for _position, entry, match in _admitted_entries(
entries, filters, links, favorites, image_tags
entries, filters, links, members, image_tags
):
if len(targets) == cap:
return ListingSelection(targets, list(positives), unlinked, True)
@@ -2382,9 +2413,11 @@ class ImageRoutesMixin:
) -> ListingSelection:
"""The images a parse_listing_filter result admits right now (#392).
Raises ValueError for a root outside the current roots, and
Raises ValueError for a root outside the current roots,
ListingSourceChanged when the filter's ``root_scope`` no longer names
its root (the sources were renumbered since the listing).
its root (the sources were renumbered since the listing), and
CollectionMissing when its ``collection`` no longer exists (checked
before any scan and again after the members are read).
"""
output_dirs = list(self._get_all_output_dirs())
root_index = _workspace_root(params, output_dirs)
@@ -2394,11 +2427,11 @@ class ImageRoutesMixin:
raise ListingSourceChanged()
raise ValueError("Invalid gallery root scope")
filters = _workspace_filters(params)
listings = [
(index, root, await self._get_gallery_files(root))
for index, root in enumerate(output_dirs)
if root_index is None or index == root_index
]
if filters["collection"] is not None and not await self._run_in_executor(
self.db.collection_exists, filters["collection"]
):
raise CollectionMissing("No such collection")
listings = await self._workspace_listings(output_dirs, root_index)
return await self._run_in_executor(
self._select_workspace_matches, listings, filters, cap
)
@@ -3583,25 +3616,58 @@ class ImageRoutesMixin:
self._favorite_listing_cache[cache_key] = (files, listing)
return listing
def _workspace_favorite_paths(self, listings: list[tuple]) -> frozenset:
"""The listed paths of every favorite in the selected roots (blocking).
def _workspace_members(self, listings: list[tuple], filters: dict):
"""The member set the filters ask for, or None (blocking)."""
collection = filters.get("collection")
if collection is not None:
members = self._workspace_member_paths(
listings,
lambda root_key: self.db.collection_members(collection, root_key),
)
# Checked after the reads: a collection deleted before or while
# they ran is a 404, never an empty page passing for it.
if not self.db.collection_exists(collection):
raise CollectionMissing("No such collection")
return members
if filters["favorites"]:
return self._workspace_favorite_paths(listings)
return None
One primary-key read per root (get_root_favorites), matched against the
root's cached ``{relative key: listing spelling}``, so the work follows
the favorites rather than the files, and a key folds case exactly as
the toggle stored it (Windows). A favorite whose file is gone matches
nothing. A failed read raises: an unfiltered page must never pass for
the favorites.
def _workspace_favorite_paths(self, listings: list[tuple]) -> frozenset:
"""The listed paths of every favorite in the selected roots (blocking)."""
return self._workspace_member_paths(listings, self.db.get_root_favorites)
def _workspace_member_paths(
self, listings: list[tuple], keys_for_source: Callable[[str], Iterable[str]]
) -> frozenset:
"""The listed paths of the stored keys of each selected root (blocking).
One read per root (``keys_for_source(root key)``: the favorites' or a
collection's primary-key range), matched against the root's cached
``{relative key: listing spelling}``, so the work follows the members
rather than the files, and a key folds case exactly as it was stored
(Windows). A member whose file is gone matches nothing. A failed read
raises: an unfiltered page must never pass for the members.
"""
paths = set()
for _index, root, files in listings:
listing = self._favorite_listing(root, files)
for key in self.db.get_root_favorites(_favorite_root_key(root)):
for key in keys_for_source(_favorite_root_key(root)):
listed = listing.get(key)
if listed is not None:
paths.add(root / listed)
return frozenset(paths)
async def _workspace_listings(
self, output_dirs: list[Path], root_index: int | None
) -> list[tuple]:
"""``(index, root, scan)`` of every selected root (cached scans)."""
return [
(index, root, await self._get_gallery_files(root))
for index, root in enumerate(output_dirs)
if root_index is None or index == root_index
]
def _workspace_favorites(self, page: list[tuple]) -> set[tuple[int, str]]:
"""``(root index, listing path)`` of the favorites among *page* entries."""
by_root: dict[tuple[int, Path], dict[str, str]] = {}
+54
View File
@@ -0,0 +1,54 @@
"""The listing identity of one media file, as a request names it.
Shared by the image-annotation routes and the collection routes (#331): a
target is ``{root_index, root_scope, path}``, checked here without touching
the filesystem or the database. The scope is compared with the current roots
by the caller.
"""
import os
import re
from .images import _forbidden, _has_traversal
from .prompts import bad_request
try:
from ...utils.media_capture import CAPTURE_EXTENSIONS
except ImportError:
from utils.media_capture import CAPTURE_EXTENSIONS
MAX_TARGET_PATH = 1024
TARGET_FIELDS = frozenset({"root_index", "root_scope", "path"})
_HEX64 = re.compile(r"[0-9a-f]{64}")
TAGGING_MEDIA_MESSAGE = "Only media files can be tagged"
def media_target(data: object, *, media_message: str = TAGGING_MEDIA_MESSAGE):
"""``((index, scope, relative), None)`` for a target, else ``(None, error)``.
Only the listing identity is accepted: an int root index, the root's
64-hex scope and a bounded root-relative POSIX media path, never an
absolute, drive-qualified, backslash or ``..`` path (a 400 for a malformed
shape, a 403 for an unsafe or non-media path, carrying *media_message*).
"""
if not isinstance(data, dict) or set(data) != TARGET_FIELDS:
return None, bad_request("target must name root_index, root_scope and path")
index, scope, relative = data["root_index"], data["root_scope"], data["path"]
if type(index) is not int or index < 0:
return None, bad_request("Invalid gallery root")
if not isinstance(scope, str) or _HEX64.fullmatch(scope) is None:
return None, bad_request("Invalid gallery root scope")
if not isinstance(relative, str) or not relative or len(relative) > MAX_TARGET_PATH:
return None, bad_request("path must be a relative media path")
parts = relative.split("/")
if (
"\0" in relative
or "\\" in relative
or relative.startswith("/")
or _has_traversal(relative)
or any(part in ("", ".", "..") or ":" in part for part in parts)
):
return None, _forbidden()
if os.path.splitext(parts[-1])[1].lower() not in CAPTURE_EXTENSIONS:
return None, _forbidden(media_message)
return (index, scope, relative), None
+4
View File
@@ -1227,7 +1227,9 @@ class PromptRoutesMixin:
"""
from .images import (
MAX_LISTING_SELECTION,
CollectionMissing,
ListingSourceChanged,
collection_not_found,
parse_listing_filter,
)
@@ -1252,6 +1254,8 @@ class PromptRoutesMixin:
},
status=409,
)
except CollectionMissing:
return collection_not_found()
except ValueError as exc:
return refused(str(exc))
if selection.truncated:
+2
View File
@@ -262,6 +262,8 @@ def create_app(root: Path) -> web.Application:
with target.db.model.get_connection() as conn:
for table in (
"saved_views",
"image_collection_items",
"image_collections",
"image_favorites",
"media_review_decisions",
"media_review_items",
+494
View File
@@ -0,0 +1,494 @@
"""Collections (#331): the routes under /prompt_manager/collections."""
import json
import ntpath
import sqlite3
from unittest.mock import patch
try:
from test_api_images import ImageAPITestCase, make_png
from open_handles import assert_closed
except ImportError:
from tests.test_api_images import ImageAPITestCase, make_png
from tests.open_handles import assert_closed
BASE = "/prompt_manager/collections"
ITEMS = BASE + "/items"
MEMBERSHIP = BASE + "/membership"
SOURCES_CHANGED = "The sources changed. Refresh the images and try again."
class CollectionRoutesTestCase(ImageAPITestCase):
async def tearDownAsync(self):
self.api.db.close_all()
assert_closed(self, self._temp_db.name, self._output_tmp.name)
await super().tearDownAsync()
def make(self, *names, root=None):
for name in names:
make_png((root or self.output_dir) / name)
def scope(self, index=0):
roots = list(self.api._get_all_output_dirs())
return self.api._gallery_root_scope(roots[index], index)
def item(self, path, index=0, scope=None):
return {
"root_index": index,
"root_scope": scope or self.scope(index),
"path": path,
}
def two_sources(self):
second = self.comfy_dir / "second"
second.mkdir()
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
return second
async def call(self, method, url, body=None, status=200, **params):
response = await self.client.request(method, url, json=body, params=params)
text = await response.text()
self.assertEqual(response.status, status, text)
data = json.loads(text)
self.assertNotIn(self.comfy_dir.as_posix(), text)
self.assertNotIn(str(self.comfy_dir).replace("\\", "\\\\"), text)
return data
async def add(self, status=200, **body):
return await self.call("POST", ITEMS, body, status)
async def listed(self):
return (await self.call("GET", BASE))["collections"]
class TestCollectionRoutes(CollectionRoutesTestCase):
async def test_each_route_succeeds(self):
self.make("a.png", "b.png", "c.png")
self.assertEqual(
await self.call("GET", BASE), {"success": True, "collections": []}
)
made = await self.add(
name=" LoRa testing ", items=[self.item("a.png"), self.item("b.png")]
)
grids = made["collection"]
self.assertEqual(grids["name"], "LoRa testing")
self.assertEqual(
made,
{
"success": True,
"collection": grids,
"created": True,
"added": 2,
"already": 0,
"not_found": 0,
},
)
again = await self.add(
collection_id=grids["id"], items=[self.item("b.png"), self.item("c.png")]
)
self.assertEqual(
(again["created"], again["added"], again["already"]), (False, 1, 1)
)
self.assertEqual(
await self.listed(),
[{"id": grids["id"], "name": "LoRa testing", "count": 3, "missing": 0}],
)
membership = await self.call(
"GET", MEMBERSHIP, root_index="0", root_scope=self.scope(), path="a.png"
)
self.assertEqual(membership, {"success": True, "collection_ids": [grids["id"]]})
removed = await self.call(
"DELETE",
ITEMS,
{"collection_id": grids["id"], "items": [self.item("a.png")]},
)
self.assertEqual(removed, {"success": True, "collection": grids, "removed": 1})
renamed = await self.call(
"PATCH", f"{BASE}/{grids['id']}", {"name": "LoRA grids"}
)
self.assertEqual(
renamed,
{"success": True, "collection": {"id": grids["id"], "name": "LoRA grids"}},
)
deleted = await self.call("DELETE", f"{BASE}/{grids['id']}")
self.assertEqual(deleted, {"success": True, "deleted": True})
self.assertEqual(await self.listed(), [])
# Files are never touched.
self.assertTrue((self.output_dir / "a.png").is_file())
async def test_names_are_unique_ignoring_case(self):
self.make("a.png", "b.png")
first = (await self.add(name="LoRa testing", items=[self.item("a.png")]))[
"collection"
]
second = await self.add(name="lora TESTING", items=[self.item("b.png")])
self.assertEqual(second["collection"], first)
self.assertFalse(second["created"])
self.assertEqual(len(await self.listed()), 1)
async def test_bad_requests_are_400(self):
self.make("a.png")
made = await self.add(name="Set", items=[self.item("a.png")])
cid = made["collection"]["id"]
good = [self.item("a.png")]
for body in (
{"items": good},
{"name": "x", "collection_id": cid, "items": good},
{"name": "x"},
{"name": "x", "items": good, "extra": 1},
{"collection_id": "1", "items": good},
{"collection_id": True, "items": good},
{"collection_id": 0, "items": good},
{"collection_id": 10**10, "items": good},
{"name": "", "items": good},
{"name": "x" * 101, "items": good},
{"name": "bell\x07", "items": good},
{"name": 5, "items": good},
{"name": "x", "items": []},
{"name": "x", "items": "a.png"},
{"name": "x", "items": [self.item("a.png")] * 501},
{"name": "x", "items": [{"root_index": 0, "path": "a.png"}]},
{"name": "x", "items": [{**self.item("a.png"), "root_index": "0"}]},
{"name": "x", "items": [{**self.item("a.png"), "root_scope": "xyz"}]},
):
with self.subTest(body=body):
data = await self.add(400, **body)
self.assertFalse(data["success"])
for body in ({"items": good}, {"collection_id": cid}):
with self.subTest(remove=body):
await self.call("DELETE", ITEMS, body, 400)
for raw in ("abc", "0", "1234567890", "-1", "1.0"):
with self.subTest(path_id=raw):
await self.call("PATCH", f"{BASE}/{raw}", {"name": "x"}, 400)
await self.call("DELETE", f"{BASE}/{raw}", None, 400)
await self.call("PATCH", f"{BASE}/{cid}", {"name": ""}, 400)
await self.call("PATCH", f"{BASE}/{cid}", {"label": "x"}, 400)
for params in (
{"root_index": "x", "root_scope": self.scope(), "path": "a.png"},
{"root_index": "0", "path": "a.png"},
{"root_index": "0", "root_scope": "nope", "path": "a.png"},
):
with self.subTest(membership=params):
await self.call("GET", MEMBERSHIP, status=400, **params)
self.assertEqual((await self.listed())[0]["count"], 1)
async def test_unsafe_paths_are_403(self):
self.make("a.png")
for path in (
"../a.png",
"sub/../../a.png",
"C:/a.png",
"a\\b.png",
"/a.png",
"notes.txt",
):
with self.subTest(path=path):
await self.add(403, name="Set", items=[self.item(path)])
await self.call(
"GET",
MEMBERSHIP,
status=403,
root_index="0",
root_scope=self.scope(),
path=path,
)
self.assertEqual(await self.listed(), [])
async def test_a_body_over_the_cap_is_413(self):
response = await self.client.post(
ITEMS, data=b'{"name":"x","items":[' + b" " * (1024 * 1024 + 1) + b"]}"
)
self.assertEqual(response.status, 413)
async def test_missing_collection_and_files_are_404(self):
self.make("a.png")
data = await self.add(404, collection_id=999, items=[self.item("a.png")])
self.assertEqual(data["code"], "collection-not-found")
data = await self.call(
"DELETE", ITEMS, {"collection_id": 999, "items": [self.item("a.png")]}, 404
)
self.assertEqual(data["code"], "collection-not-found")
data = await self.call("PATCH", f"{BASE}/999", {"name": "x"}, 404)
self.assertEqual(data["code"], "collection-not-found")
data = await self.call("DELETE", f"{BASE}/999", None, 404)
self.assertEqual(data["code"], "collection-not-found")
data = await self.add(404, name="New", items=[self.item("never.png")])
self.assertEqual(data["code"], "files-not-found")
self.assertEqual(await self.listed(), [])
async def test_unlisted_files_count_as_not_found(self):
self.make("a.png")
made = await self.add(
name="Set", items=[self.item("a.png"), self.item("later.png")]
)
self.assertEqual((made["added"], made["not_found"]), (1, 1))
again = await self.add(
collection_id=made["collection"]["id"],
items=[self.item("a.png"), self.item("later.png")],
)
self.assertEqual(
(again["added"], again["already"], again["not_found"]), (0, 1, 1)
)
async def test_name_taken_is_409(self):
self.make("a.png")
first = (await self.add(name="First", items=[self.item("a.png")]))["collection"]
second = (await self.add(name="Second", items=[self.item("a.png")]))[
"collection"
]
data = await self.call(
"PATCH", f"{BASE}/{second['id']}", {"name": "FIRST"}, 409
)
self.assertEqual(
(data["code"], data["collection_id"]), ("name-taken", first["id"])
)
renamed = await self.call("PATCH", f"{BASE}/{second['id']}", {"name": "SECOND"})
self.assertEqual(renamed["collection"]["name"], "SECOND")
async def test_stale_scope_is_409_after_reordering_sources(self):
second = self.two_sources()
self.make("a.png")
self.make("b.png", root=second)
stale = self.item("b.png", 1)
made = await self.add(name="Set", items=[self.item("a.png")])
cid = made["collection"]["id"]
self.api._get_all_output_dirs = lambda: [second, self.output_dir]
for method, body in (
("POST", {"collection_id": cid, "items": [self.item("b.png", 0), stale]}),
("DELETE", {"collection_id": cid, "items": [stale]}),
):
with self.subTest(method=method):
data = await self.call(method, ITEMS, body, 409)
self.assertEqual(data["code"], "source-changed")
self.assertEqual(data["error"], SOURCES_CHANGED)
data = await self.call(
"GET",
MEMBERSHIP,
status=409,
root_index="1",
root_scope=stale["root_scope"],
path="b.png",
)
self.assertEqual(data["code"], "source-changed")
# Nothing was written by the refused add; the member follows its folder.
listed = await self.listed()
self.assertEqual((listed[0]["count"], listed[0]["missing"]), (1, 0))
membership = await self.call(
"GET", MEMBERSHIP, root_index="1", root_scope=self.scope(1), path="a.png"
)
self.assertEqual(membership["collection_ids"], [cid])
async def test_counts_exclude_a_missing_member(self):
second = self.two_sources()
self.make("a.png", "b.png")
self.make("c.png", root=second)
made = await self.add(
name="Grids",
items=[self.item("a.png"), self.item("b.png"), self.item("c.png", 1)],
)
other = await self.add(name="Other", items=[self.item("c.png", 1)])
(self.output_dir / "b.png").unlink()
self.api.invalidate_gallery_cache()
listed = await self.listed()
self.assertEqual(
listed,
[
{
"id": made["collection"]["id"],
"name": "Grids",
"count": 2,
"missing": 1,
},
{
"id": other["collection"]["id"],
"name": "Other",
"count": 1,
"missing": 0,
},
],
)
# The member is removable although it is no longer listed.
removed = await self.call(
"DELETE",
ITEMS,
{"collection_id": made["collection"]["id"], "items": [self.item("b.png")]},
)
self.assertEqual(removed["removed"], 1)
async def test_counts_are_null_when_a_source_scan_fails(self):
self.make("a.png")
await self.add(name="Set", items=[self.item("a.png")])
async def broken(_root):
raise OSError("scan failed at /secret/place")
with patch.object(self.api, "_get_gallery_files", side_effect=broken):
data = await self.call("GET", BASE)
self.assertEqual(
data["collections"][0] | {"id": 0},
{"id": 0, "name": "Set", "count": None, "missing": None},
)
self.assertNotIn("secret", json.dumps(data))
async def test_windows_case_spellings_are_one_member(self):
self.make("Sub/A.png")
with patch("py.api.images._path_flavor", ntpath):
made = await self.add(
name="Set", items=[self.item("Sub/A.png"), self.item("SUB/a.PNG")]
)
self.assertEqual((made["added"], made["already"]), (1, 0))
(self.output_dir / "Sub" / "A.png").rename(
self.output_dir / "Sub" / "a.png"
)
(self.output_dir / "Sub").rename(self.output_dir / "sub")
self.api.invalidate_gallery_cache()
listed = await self.listed()
self.assertEqual((listed[0]["count"], listed[0]["missing"]), (1, 0))
membership = await self.call(
"GET",
MEMBERSHIP,
root_index="0",
root_scope=self.scope(),
path="sub/a.png",
)
self.assertEqual(membership["collection_ids"], [made["collection"]["id"]])
response = await self.client.get(
"/prompt_manager/images/output",
params={
"mode": "workspace",
"collection": str(made["collection"]["id"]),
},
)
page = await response.json()
self.assertEqual([i["path"] for i in page["images"]], ["sub/a.png"])
async def test_posix_case_distinct_files_stay_two_members(self):
self.make("A.png", "a.png")
if len(list(self.output_dir.iterdir())) != 2:
self.skipTest("case-insensitive filesystem")
made = await self.add(
name="Set", items=[self.item("A.png"), self.item("a.png")]
)
self.assertEqual(made["added"], 2)
async def test_a_500_never_carries_the_exception_text(self):
self.make("a.png")
made = await self.add(name="Set", items=[self.item("a.png")])
cid = made["collection"]["id"]
def broken(*_args, **_kwargs):
raise sqlite3.OperationalError("disk I/O error at /secret/prompts.db")
cases = (
(
"list_collections",
"GET",
BASE,
None,
{},
"Collections could not be read.",
),
(
"file_collections",
"GET",
MEMBERSHIP,
None,
{"root_index": "0", "root_scope": self.scope(), "path": "a.png"},
"Collections could not be read.",
),
(
"add_to_collection",
"POST",
ITEMS,
{"collection_id": cid, "items": [self.item("a.png")]},
{},
"The collection could not be changed.",
),
(
"remove_from_collection",
"DELETE",
ITEMS,
{"collection_id": cid, "items": [self.item("a.png")]},
{},
"The collection could not be changed.",
),
(
"rename_collection",
"PATCH",
f"{BASE}/{cid}",
{"name": "x"},
{},
"The collection could not be changed.",
),
(
"delete_collection",
"DELETE",
f"{BASE}/{cid}",
None,
{},
"The collection could not be changed.",
),
)
for method_name, method, url, body, params, message in cases:
with self.subTest(method=method_name):
with patch.object(self.api.db, method_name, side_effect=broken):
data = await self.call(method, url, body, 500, **params)
self.assertEqual(data, {"success": False, "error": message})
class TestCollectionListingFilter(CollectionRoutesTestCase):
"""#392's "all N matching" filter carries ``collection`` like the listing."""
async def test_a_filter_with_a_collection_resolves_only_members(self):
second = self.two_sources()
self.make("a.png", "b.png")
self.make("c.png", root=second)
made = await self.add(
name="Set", items=[self.item("a.png"), self.item("c.png", 1)]
)
cid = str(made["collection"]["id"])
selection = await self.api.resolve_listing_filter({"collection": cid})
self.assertEqual(
sorted((index, path) for index, _scope, path in selection.targets),
[(0, "a.png"), (1, "c.png")],
)
selection = await self.api.resolve_listing_filter(
{"collection": cid, "root": "1", "root_scope": self.scope(1)}
)
self.assertEqual([t[2] for t in selection.targets], ["c.png"])
for bad in ("0", "x", "1234567890"):
with self.subTest(bad=bad):
with self.assertRaises(ValueError):
from py.api.images import parse_listing_filter
parse_listing_filter({"collection": bad})
async def test_bulk_category_acts_on_collection_members_only(self):
self.make("a.png", "b.png")
inside = self._save_prompt("Inside", category="Old")
outside = self._save_prompt("Outside", category="Old")
self.api.db.link_image_to_prompt(inside, str(self.output_dir / "a.png"))
self.api.db.link_image_to_prompt(outside, str(self.output_dir / "b.png"))
made = await self.add(name="Set", items=[self.item("a.png")])
cid = str(made["collection"]["id"])
data = await self.call(
"POST",
"/prompt_manager/bulk/category",
{"images_filter": {"collection": cid}, "category": "New"},
)
self.assertEqual((data["updated_count"], data["prompt_count"]), (1, 1))
self.assertEqual(self.api.db.get_prompt_by_id(inside)["category"], "New")
self.assertEqual(self.api.db.get_prompt_by_id(outside)["category"], "Old")
async def test_a_deleted_collection_in_a_bulk_filter_is_a_404(self):
self.make("a.png")
data = await self.call(
"POST",
"/prompt_manager/bulk/category",
{"images_filter": {"collection": "77"}, "category": "New"},
status=404,
)
self.assertEqual(data["code"], "collection-not-found")
with self.assertRaises(Exception) as caught:
await self.api.resolve_listing_filter({"collection": "77"})
self.assertEqual(type(caught.exception).__name__, "CollectionMissing")
+168
View File
@@ -2494,3 +2494,171 @@ class TestListingFilterSelection(ImageAPITestCase):
self.assertEqual(response.status, 409)
self.assertEqual((await response.json())["code"], "source-changed")
class TestWorkspaceCollectionFilter(ImageAPITestCase):
"""#331: ``collection=<id>`` keeps only that collection's listed members."""
media = TestWorkspaceImages.media
get = TestWorkspaceImages.get
_names = TestWorkspaceImages._names
async def tearDownAsync(self):
self.api.db.close_all()
assert_closed(self, self._temp_db.name, self._output_tmp.name)
await super().tearDownAsync()
def keys(self, root, relative):
from py.api.images import _favorite_relative_key, _favorite_root_key
return (_favorite_root_key(root), _favorite_relative_key(relative))
def collect(self, name, *members):
"""A collection holding ``(root, relative)`` members; its id."""
result = self.api.db.add_to_collection(
name=name, items=[self.keys(root, rel) for root, rel in members]
)
return result["collection"]["id"]
def two_sources(self):
second = self.comfy_dir / "second"
second.mkdir()
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
return second
async def test_listing_keeps_only_collection_members(self):
second = self.two_sources()
self.media("a.png", 300)
self.media("b.png", 200)
self.media("sub/c.png", 100, root=second)
self.media("d.png", 50, root=second)
grids = self.collect(
"Grids",
(self.output_dir, "a.png"),
(second, "sub/c.png"),
(self.output_dir, "gone.png"), # stored, never listed
)
self.collect("Other", (self.output_dir, "b.png"))
response, data = await self.get(collection=str(grids))
self.assertEqual(response.status, 200, data)
self.assertEqual(self._names(data), ["a.png", "c.png"])
self.assertEqual(data["total"], 2)
_, data = await self.get(collection=str(grids), root="1")
self.assertEqual(self._names(data), ["c.png"])
# Without the filter the listing is unchanged.
_, data = await self.get()
self.assertEqual(data["total"], 4)
async def test_collection_filter_combines_with_the_other_filters(self):
night = self._save_prompt("Neon zebra", category="Night", tags=["a"])
day = self._save_prompt("Morning field", category="Day", tags=["a"])
files = {
name: self.media(name, 100 + index)
for index, name in enumerate(
("night.png", "day.png", "zebra_plain.png", "clip.mp4", "out.png")
)
}
self.api.db.link_image_to_prompt(night, str(files["night.png"]))
self.api.db.link_image_to_prompt(day, str(files["day.png"]))
members = [(self.output_dir, name) for name in files if name != "out.png"]
grids = str(self.collect("Grids", *members))
cases = (
({}, ["clip.mp4", "day.png", "night.png", "zebra_plain.png"]),
({"q": "zebra"}, ["night.png", "zebra_plain.png"]),
({"category": "Night"}, ["night.png"]),
({"tags": "a"}, ["day.png", "night.png"]),
({"media": "video"}, ["clip.mp4"]),
({"media": "image", "q": "day"}, ["day.png"]),
)
for params, expected in cases:
with self.subTest(**params):
response, data = await self.get(collection=grids, sort="name", **params)
self.assertEqual(response.status, 200, data)
self.assertEqual(self._names(data), expected)
self.assertEqual(data["total"], len(expected))
async def test_invalid_collection_is_a_400(self):
for value in ("x", "0", "-1", "1.5", "1234567890", " 1"):
with self.subTest(value=value):
response, data = await self.get(collection=value)
self.assertEqual(response.status, 400, data)
async def test_unknown_collection_is_a_404(self):
self.media("a.png")
response, data = await self.get(collection="77")
self.assertEqual(response.status, 404)
self.assertEqual(data["code"], "collection-not-found")
self.assertNotIn("images", data)
self.api._get_all_output_dirs = lambda: []
response, data = await self.get(collection="77")
self.assertEqual((response.status, data["code"]), (404, "collection-not-found"))
async def test_a_collection_deleted_during_the_read_is_a_404(self):
self.media("a.png")
grids = self.collect("Grids", (self.output_dir, "a.png"))
real = self.api.db.collection_members
def deleted_meanwhile(collection_id, root_key):
self.api.db.delete_collection(collection_id)
return real(collection_id, root_key)
with patch.object(
self.api.db, "collection_members", side_effect=deleted_meanwhile
):
response, data = await self.get(collection=str(grids))
self.assertEqual((response.status, data["code"]), (404, "collection-not-found"))
async def test_a_failed_members_read_is_a_500_never_an_unfiltered_page(self):
self.media("a.png")
self.media("b.png")
grids = self.collect("Grids", (self.output_dir, "a.png"))
def broken(*_args, **_kwargs):
raise sqlite3.OperationalError("database is locked")
with patch.object(self.api.db, "collection_members", side_effect=broken):
response, data = await self.get(collection=str(grids))
self.assertEqual(response.status, 500)
self.assertNotIn("images", data)
self.assertNotIn("locked", json.dumps(data))
async def test_reordering_sources_keeps_members(self):
second = self.two_sources()
self.media("a.png", root=second)
self.media("b.png")
grids = str(self.collect("Grids", (second, "a.png")))
_, data = await self.get(collection=grids)
self.assertEqual(
[(i["root_index"], i["path"]) for i in data["images"]], [(1, "a.png")]
)
self.api._get_all_output_dirs = lambda: [second, self.output_dir]
_, data = await self.get(collection=grids)
self.assertEqual(
[(i["root_index"], i["path"]) for i in data["images"]], [(0, "a.png")]
)
async def test_cursor_from_another_collection_is_refused(self):
for index, name in enumerate(("a.png", "b.png", "c.png")):
self.media(name, 100 + index)
everything = [(self.output_dir, name) for name in ("a.png", "b.png", "c.png")]
first = str(self.collect("First", *everything))
second = str(self.collect("Second", *everything))
_, page = await self.get(collection=first, limit=1)
token = page["next_cursor"]
self.assertIsInstance(token, str)
response, _ = await self.get(collection=first, limit=1, cursor=token)
self.assertEqual(response.status, 200)
for params in ({"collection": second}, {}, {"favorites": "1"}):
with self.subTest(**params):
response, _ = await self.get(cursor=token, limit=1, **params)
self.assertEqual(response.status, 400)
async def test_collection_wins_over_favorites(self):
self.media("a.png")
self.media("b.png")
grids = str(self.collect("Grids", (self.output_dir, "a.png")))
from py.api.images import _favorite_root_key
self.api.db.set_image_favorite(_favorite_root_key(self.output_dir), "b.png")
_, data = await self.get(collection=grids, favorites="1")
self.assertEqual(self._names(data), ["a.png"])
+10
View File
@@ -884,6 +884,16 @@ class TestFilterStart(ImageTagJobTestCase):
self.assertEqual((status, data["code"]), (400, "no-images"))
self.assertEqual(self.api._jobs.list(), [])
async def test_filter_with_a_deleted_collection_is_a_404_without_a_job(self):
self.png("a.png")
status, data = await self.start(
{"mode": "add", "filter": {"collection": "77"}, "tags": ["x"]}
)
self.assertEqual((status, data["code"]), (404, "collection-not-found"))
self.assertEqual(self.api._jobs.list(), [])
async def test_filter_over_the_cap_is_refused_without_a_job(self):
from py.api import image_tag_jobs