server+client: use safety rating param in "lookalikes" api

3 files changed, 79 insertions, 10 deletions

diff --git a/server/szurubooru/api/post_api.py b/server/szurubooru/api/post_api.py
index e5ca037..4a629fe 100644
--- a/server/szurubooru/api/post_api.py
+++ b/server/szurubooru/api/post_api.py
@@ -328,6 +328,7 @@ def get_posts_lookalikes(
     auth.verify_privilege(ctx.user, "posts:reverse_search")
     limit = ctx.get_param_as_int("limit", default=10, min=1, max=100)
     threshold = ctx.get_param_as_float("threshold", default=1, min=0, max=100)
+    query_text = ctx.get_param_as_string("query", default="")
     post_id = _get_post_id(params)
     post = posts.get_post_by_id(post_id)
     if post.signature is None:
@@ -335,7 +336,7 @@ def get_posts_lookalikes(
 
     sig = image_hash.unpack_signature(post.signature.signature)
     # limit + 1 because the original post will be excluded
-    lookalikes = posts.search_by_signature(sig, limit + 1, threshold)
+    lookalikes = posts.search_by_signature(sig, limit + 1, threshold, query_text)
     # exclude the original post:
     lookalikes = filter(lambda la: la[1].post_id != post_id, lookalikes)
     lookalikes = sorted(lookalikes, key=lambda la: la[0])
diff --git a/server/szurubooru/func/posts.py b/server/szurubooru/func/posts.py
index fac7cd5..a7c111a 100644
--- a/server/szurubooru/func/posts.py
+++ b/server/szurubooru/func/posts.py
@@ -1,5 +1,6 @@
 import hmac
 import logging
+import re
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
@@ -22,6 +23,7 @@ from szurubooru.func import (
     util,
 )
 from szurubooru.func.image_hash import NpMatrix
+from szurubooru.search import parser, criteria
 
 logger = logging.getLogger(__name__)
 
@@ -962,6 +964,7 @@ def search_by_signature(
     signature: NpMatrix,
     limit: int = 100,
     distance_cutoff: float = image_hash.DISTANCE_CUTOFF,
+    query_text: str = ''
 ) -> List[Tuple[float, model.Post]]:
     query_words = image_hash.generate_words(signature)
     """
@@ -972,15 +975,33 @@ def search_by_signature(
     https://www.postgresql.org/docs/9.2/functions-array.html
     """
 
-    dbquery = """
-    SELECT s.post_id, s.signature, count(a.query) AS score
-    FROM post_signature AS s, unnest(s.words, :q) AS a(word, query)
-    WHERE a.word = a.query
-    GROUP BY s.post_id
-    ORDER BY score DESC LIMIT :limit;
-    """
+    # optimization: don't join if safety is not queried:
+    if len(query_text) > 0:
+        dbquery = """
+        SELECT s.post_id, s.signature, count(a.query) AS score
+        FROM post_signature AS s
+        CROSS JOIN unnest(s.words, :q) AS a(word, query)
+        INNER JOIN post ON post.id = s.post_id
+        WHERE a.word = a.query
+        AND post.safety in :safety
+        GROUP BY s.post_id
+        ORDER BY score DESC LIMIT :limit;
+        """
+    else:
+        dbquery = """
+        SELECT s.post_id, s.signature, count(a.query) AS score
+        FROM post_signature AS s, unnest(s.words, :q) AS a(word, query)
+        WHERE a.word = a.query
+        GROUP BY s.post_id
+        ORDER BY score DESC LIMIT :limit;
+        """
+    allowed_rating = _get_safety_list(query_text)
 
-    candidates = db.session.execute(dbquery, {"q": query_words, "limit": limit})
+    candidates = db.session.execute(dbquery, {
+        "q": query_words,
+        "limit": limit,
+        "safety": tuple(allowed_rating),
+    })
     data = tuple(
         zip(
             *[
@@ -1001,3 +1022,41 @@ def search_by_signature(
         ]
     else:
         return []
+
+
+def _get_safety_list(query_text: str = '') -> List[str]:
+    """Will output a list of safety options matched by the query"""
+    # TODO(hunternif): searching by signature should be done in executor,
+    # together with all other tokens, but as a quick fix for safety rating,
+    # we can parse it here.
+    # Assuming format: -rating:safe,sketchy,unsafe
+    query_parser = parser.Parser()
+    search_query = query_parser.parse(query_text)
+    safety_map = util.flip(SAFETY_MAP)
+    allowed = []
+    disallowed = []
+
+    def process_safety(safety_value: str):
+        safety = safety_map.get(safety_value, None)
+        if safety:
+            if token.negated:
+                disallowed.append(safety)
+            else:
+                allowed.append(safety)
+
+    for token in search_query.named_tokens:
+        if token.name == "rating":
+            criterion = token.criterion
+            if isinstance(criterion, criteria.PlainCriterion):
+                process_safety(criterion.value)
+            elif isinstance(criterion, criteria.ArrayCriterion):
+                for value in criterion.values:
+                    process_safety(value)
+
+    if len(allowed) == 0:
+        allowed = [
+            model.Post.SAFETY_SAFE,
+            model.Post.SAFETY_SKETCHY,
+            model.Post.SAFETY_UNSAFE,
+        ]
+    return [x for x in allowed if x not in disallowed]
\ No newline at end of file
diff --git a/server/szurubooru/tests/func/test_posts.py b/server/szurubooru/tests/func/test_posts.py
index e1be764..168a74a 100644
--- a/server/szurubooru/tests/func/test_posts.py
+++ b/server/szurubooru/tests/func/test_posts.py
@@ -8,13 +8,13 @@ from szurubooru import db, model
 from szurubooru.func import (
     comments,
     files,
-    image_hash,
     images,
     posts,
     tags,
     users,
     util,
 )
+from szurubooru.func.posts import _get_safety_list
 
 
 @pytest.mark.parametrize(
@@ -1253,3 +1253,12 @@ def test_search_by_image(post_factory, config_injector, read_asset):
 
     result2 = posts.search_by_image(read_asset("png.png"))
     assert not result2
+
+
+def test_get_safety_list():
+    assert _get_safety_list('') == ['safe', 'sketchy', 'unsafe']
+    assert _get_safety_list('abc') == ['safe', 'sketchy', 'unsafe']
+    assert _get_safety_list('abc rating:lol -def') ==\
+           ['safe', 'sketchy', 'unsafe']
+    assert _get_safety_list('abc -rating:sketchy,lol def') == ['safe', 'unsafe']
+    assert _get_safety_list('rating:safe,unsafe -rating:safe') == ['unsafe']