Significant performance improvement across tagGalleriesFromImages

plugins/tagGalleriesFromImages/tagGalleriesFromImages.py

@@ -2,13 +2,22 @@
 from stashapi.stashapp import StashInterface
 import sys
 import json
+import time
+
+GALLERY_PAGE_SIZE = 150
+tag_cache = {}
+
 
 def processAll():
     exclusion_marker_tag_id = None
-    if settings["excludeWithTag"] != "":
-        exclussion_marker_tag = stash.find_tag(settings["excludeWithTag"])
-        if exclussion_marker_tag is not None:
-            exclusion_marker_tag_id = exclussion_marker_tag['id']
+    if settings["excludeWithTag"]:
+        if settings["excludeWithTag"] in tag_cache:
+            exclusion_marker_tag_id = tag_cache[settings["excludeWithTag"]]['id']
+        else:
+            exclusion_marker_tag = stash.find_tag(settings["excludeWithTag"])
+            if exclusion_marker_tag:
+                exclusion_marker_tag_id = exclusion_marker_tag['id']
+                tag_cache[settings["excludeWithTag"]] = exclusion_marker_tag
 
     query = {
         "image_count": {
@@ -18,95 +27,137 @@ def processAll():
     }
     if settings['excludeOrganized']:
         query["organized"] = False
-    if exclusion_marker_tag_id is not None:
+    if exclusion_marker_tag_id:
         query["tags"] = {
             "value": [exclusion_marker_tag_id],
             "modifier": "EXCLUDES"
         }
 
-    total_count = stash.find_galleries(f=query, filter={"page": 0, "per_page": 0}, get_count=True)[0]
-    i = 0
-    while i < total_count:
-        log.progress((i / total_count))
+    try:
+        total_count = stash.find_galleries(f=query, filter={"page": 1, "per_page": 1}, get_count=True)[0]
+    except Exception:
+        total_count = 0
+
+    processed = 0
+    page = 1
+
+    while True:
+        if total_count > 0:
+            log.progress(min(processed / total_count, 1.0))
+
+        galleries = stash.find_galleries(
+            f=query, 
+            filter={"page": page, "per_page": GALLERY_PAGE_SIZE},
+            fragment="id title code organized tags { id name } performers { id } studio { id }"
+        )
 
-        galleries = stash.find_galleries(f=query, filter={"page": i, "per_page": 1})
-        if len(galleries) == 0:
+        if not galleries:
+            log.info("Finished processing all galleries.")
             break
-        gallery = galleries[0]
 
-        processGallery(gallery)
+        for gallery in galleries:
+            processGallery(gallery)
+            processed += 1
+
+        page += 1
+
+    tag_cache.clear()
+
+
+def should_exclude_gallery(gallery, exclusion_tag_name):
+    if exclusion_tag_name:
+        for tag in gallery.get("tags", []):
+            if tag["name"] == exclusion_tag_name:
+                return True
+
+    if settings['excludeOrganized'] and gallery.get('organized'):
+        return True
+
+    return False
 
-        i = i + 1
 
+def processGallery(gallery: dict):
+    if should_exclude_gallery(gallery, settings["excludeWithTag"]):
+        return
 
-def processGallery(gallery : dict):
-    tags = []
-    performersIds = []
-    should_tag = True
-    if settings["excludeWithTag"] != "":
-        for tag in gallery["tags"]:
-            if tag["name"] == settings["excludeWithTag"]:
-                should_tag = False
-                break
+    gallery_id = gallery['id']
+
+    images = stash.find_gallery_images(
+        gallery_id, 
+        fragment='tags { id } performers { id } studio { id }'
+    )
 
-    if settings['excludeOrganized']:
-        if gallery['organized']:
-            should_tag = False
-
-    if should_tag:
-        existing_tag_ids = {t['id'] for t in gallery['tags']}
-        existing_performer_ids = {p['id'] for p in gallery['performers']}
-
-        images = stash.find_gallery_images(gallery['id'], fragment='tags { id name } performers { id name }')
-        if len(images) > 0:
-            tag_ids = set()
-            tag_names = set()
-
-            performer_ids = set()
-            performer_names = set()
-
-            for image in images:
-                image_tag_ids = [tag['id'] for tag in image['tags']]
-                image_tag_names = [tag['name'] for tag in image['tags']]
-                tag_ids.update(image_tag_ids)
-                tag_names.update(image_tag_names)
-
-                image_performer_ids = [performer['id'] for performer in image['performers']]
-                image_performer_names = [performer['name'] for performer in image['performers']]
-                performer_ids.update(image_performer_ids)
-                performer_names.update(image_performer_names)
-
-            new_tags_ids = tag_ids - existing_tag_ids
-            new_performer_ids = performer_ids - existing_performer_ids
-
-            if len(new_tags_ids) > 0 or len(new_performer_ids) > 0:
-                log.info(f"updating gallery {gallery['id']} from {len(images)} images with tags {tag_names} ({len(new_tags_ids)} new) and performers {performer_names} ({len(new_performer_ids)} new)")
-                stash.update_galleries({"ids": gallery['id'], "tag_ids": {"mode": "ADD", "ids": list(new_tags_ids)}, "performer_ids": {"mode": "ADD", "ids": list(new_performer_ids)}})
+    if not images:
+        return
+
+    all_found_tag_ids = {t['id'] for img in images for t in img.get('tags', []) if 'id' in t}
+    all_found_performer_ids = {p['id'] for img in images for p in img.get('performers', []) if 'id' in p}
+    all_found_studio_ids = {img['studio']['id'] for img in images if img.get('studio') and img['studio'].get('id')}
+
+    if not all_found_tag_ids and not all_found_performer_ids and not all_found_studio_ids:
+        return
+
+    existing_tag_ids = {t['id'] for t in gallery.get('tags', [])}
+    existing_perf_ids = {p['id'] for p in gallery.get('performers', [])}
+    existing_studio_id = gallery.get('studio', {}).get('id') if gallery.get('studio') else None
+
+    missing_tags = all_found_tag_ids - existing_tag_ids
+    missing_perfs = all_found_performer_ids - existing_perf_ids
+
+    target_studio_id = list(all_found_studio_ids)[0] if all_found_studio_ids else None
+    studio_missing = bool(target_studio_id and (target_studio_id != existing_studio_id))
+
+    if not missing_tags and not missing_perfs and not studio_missing:
+        return
+
+    update_data = {"ids": [gallery_id]}
+    if missing_tags and all_found_tag_ids:
+        update_data["tag_ids"] = {"mode": "ADD", "ids": list(all_found_tag_ids)}
+    if missing_perfs and all_found_performer_ids:
+        update_data["performer_ids"] = {"mode": "ADD", "ids": list(all_found_performer_ids)}
+    if studio_missing:
+        update_data["studio_id"] = target_studio_id
+
+    gallery_name = gallery.get("title") or gallery.get("code") or f"ID {gallery_id}"
+    log.info(f"Syncing gallery '{gallery_name}' upward with structural image metadata updates.")
+
+    try:
+        stash.update_galleries(update_data)
+    except Exception as e:
+        log.error(f"Error updating gallery {gallery_id}: {str(e)}")
 
 
 json_input = json.loads(sys.stdin.read())
 FRAGMENT_SERVER = json_input["server_connection"]
 stash = StashInterface(FRAGMENT_SERVER)
+
 config = stash.get_configuration()
-settings = {
-    "excludeWithTag": "",
-    "excludeOrganized": False
-}
-if "tagGalleriesFromImages" in config["plugins"]:
+settings = {"excludeWithTag": "", "excludeOrganized": False}
+
+if config and "plugins" in config and "tagGalleriesFromImages" in config["plugins"]:
     settings.update(config["plugins"]["tagGalleriesFromImages"])
 
 if "mode" in json_input["args"]:
-    PLUGIN_ARGS = json_input["args"]["mode"]
-    if "processAll" in PLUGIN_ARGS:
+    if "processAll" in json_input["args"]["mode"]:
         processAll()
 elif "hookContext" in json_input["args"]:
-    id = json_input["args"]["hookContext"]['id']
-    if (
-        (
-            json_input["args"]["hookContext"]["type"] == "Gallery.Update.Post"
-                or json_input["args"]["hookContext"]["type"] == "Gallery.Create.Post"
-        ) and "inputFields" in json_input["args"]["hookContext"]
-        and len(json_input["args"]["hookContext"]["inputFields"]) > 2
-    ):
-        gallery = stash.find_gallery(id)
-        processGallery(gallery)
+    time.sleep(0.05)
+    hook = json_input["args"]["hookContext"]
+    hook_type = hook.get("type", "")
+
+    # --- GALLERY HOOK HANDLER ---
+    if hook_type in ["Gallery.Update.Post", "Gallery.Create.Post"]:
+        gallery = stash.find_gallery(hook['id'], fragment="id title code organized tags { id name } performers { id } studio { id }")
+        if gallery:
+            processGallery(gallery)
+
+    # --- IMAGE HOOK HANDLER ---
+    elif hook_type in ["Image.Update.Post", "Image.Create.Post"]:
+        image_id = hook['id']
+        image_data = stash.find_image(image_id, fragment="gallery { id }")
+
+        if image_data and image_data.get("gallery"):
+            gallery_id = image_data["gallery"]["id"]
+            gallery = stash.find_gallery(gallery_id, fragment="id title code organized tags { id name } performers { id } studio { id }")
+            if gallery:
+                processGallery(gallery)

plugins/tagGalleriesFromImages/tagGalleriesFromImages.yml

@@ -1,14 +1,14 @@
 name: Tag galleries from images
-description: tags galleries with tags of contained images.
-version: 0.1
+description: Tags galleries with tags of contained images.
+version: 1.0
 url: https://discourse.stashapp.cc/t/tag-galleries-from-images/3904
 exec:
   - python
   - "{pluginDir}/tagGalleriesFromImages.py"
 interface: raw
 
 hooks:
-  - name: update gallery
+  - name: Update gallery
     description: Will tag galleries with tags of contained images
     triggeredBy:
       - Gallery.Update.Post
@@ -17,15 +17,15 @@ hooks:
 settings:
   excludeOrganized:
     displayName: Exclude galleries marked as organized
-    description: Do not automatically tag galleries if it is marked as organized
+    description: Do not automatically tag galleries if it is marked as organized.
     type: BOOLEAN
   excludeWithTag:
     displayName: Exclude galleries with tag from Hook
-    description: Do not automatically tag galleries if the gallery has this tag
+    description: Do not automatically tag galleries if the gallery has this tag.
     type: STRING
 
 tasks:
   - name: "Tag all galleries"
-    description: Loops through all galleries, and applies the tags of the contained images. Can take a long time on large db's.
+    description: Loops through all galleries, and applies the tags of the contained images. Can take a long time on large databases.
     defaultArgs:
       mode: processAll