server: implement search for similar posts based on tags

3 files changed, 102 insertions, 1 deletions

diff --git a/server/szurubooru/api/post_api.py b/server/szurubooru/api/post_api.py
index a83ccc0..edb78fb 100644
--- a/server/szurubooru/api/post_api.py
+++ b/server/szurubooru/api/post_api.py
@@ -11,6 +11,7 @@ from szurubooru.func import (
     posts,
     scores,
     serialization,
+    similar,
     snapshots,
     tags,
     versions,
@@ -322,7 +323,8 @@ def get_posts_by_image(
 
 @rest.routes.get("/posts/median/?")
 def get_posts_median(
-        ctx: rest.Context, _params: Dict[str, str] = {}) -> rest.Response:
+    ctx: rest.Context, _params: Dict[str, str] = {}
+) -> rest.Response:
     auth.verify_privilege(ctx.user, "posts:list")
     _search_executor_config.user = ctx.user
     query_text = ctx.get_param_as_string("query", default="")
@@ -336,3 +338,21 @@ def get_posts_median(
         "total": len(results),
         "results": list([_serialize_post(ctx, post) for post in results])
     }
+
+
+@rest.routes.get("/post/(?P<post_id>[^/]+)/similar/?")
+def get_posts_similar(
+    ctx: rest.Context, params: Dict[str, str]
+) -> rest.Response:
+    auth.verify_privilege(ctx.user, "posts:view:similar")
+    _search_executor_config.user = ctx.user
+    post_id = _get_post_id(params)
+    post = posts.get_post_by_id(post_id)
+    limit = ctx.get_param_as_int("limit", default=10, min=1, max=100)
+    results = similar.find_similar_posts(post, limit)
+    return {
+        "limit": limit,
+        "results": list([
+            posts.serialize_micro_post(result, ctx.user) for result in results
+        ])
+    }
diff --git a/server/szurubooru/func/similar.py b/server/szurubooru/func/similar.py
new file mode 100644
index 0000000..3955c03
--- /dev/null
+++ b/server/szurubooru/func/similar.py
@@ -0,0 +1,38 @@
+from queue import Queue
+from typing import List
+
+from szurubooru import model, search
+
+_search_executor_config = search.configs.PostSearchConfig()
+_search_executor = search.Executor(_search_executor_config)
+
+
+def find_similar_posts(source_post: model.Post, limit: int) -> List[model.Post]:
+    results = []
+    queue = Queue()  # contains lists of tags to search
+    queue.put(source_post.tags)
+
+    while not queue.empty():
+        # put follow-up searches on the queue
+        last_tags = queue.get()
+        if len(last_tags) > 1:
+            for removed_tag in last_tags:
+                next_search = list(filter(lambda t: t != removed_tag, last_tags))
+                queue.put(next_search)
+
+        # prepare the current search, remove known results
+        query = ' '.join([t.first_name for t in last_tags])
+        query += ' -id:%d' % source_post.post_id
+        for r in results:
+            query += ' -id:%d' % r.post_id
+
+        # execute
+        _, posts = _search_executor.execute(query, 0, limit)
+
+        # update results
+        for p in posts:
+            results.append(p)
+            if len(results) >= limit:
+                return results
+
+    return results
diff --git a/server/szurubooru/tests/func/test_similar.py b/server/szurubooru/tests/func/test_similar.py
new file mode 100644
index 0000000..1a147db
--- /dev/null
+++ b/server/szurubooru/tests/func/test_similar.py
@@ -0,0 +1,43 @@
+import pytest
+from szurubooru import db
+from szurubooru.func import similar
+
+
+@pytest.fixture
+def verify_posts():
+    def verify(actual_posts, expected_posts):
+        actual_post_ids = list([p.post_id for p in actual_posts])
+        expected_post_ids = list([p.post_id for p in expected_posts])
+        assert actual_post_ids == expected_post_ids
+    return verify
+
+
+def test_find_similar_posts(post_factory, tag_factory, verify_posts):
+    tagA = tag_factory(names=["a"])
+    tagB = tag_factory(names=["b"])
+    tagC = tag_factory(names=["c"])
+    postA = post_factory(id=1, tags=[tagA])
+    postAB = post_factory(id=2, tags=[tagA, tagB])
+    postAC = post_factory(id=3, tags=[tagA, tagC])
+    postABC = post_factory(id=4, tags=[tagA, tagB, tagC])
+    postBC = post_factory(id=5, tags=[tagB, tagC])
+    db.session.add_all([tagA, tagB, tagC, postA, postAB, postAC, postABC, postBC])
+    db.session.flush()
+
+    results = similar.find_similar_posts(postBC, 10)
+    verify_posts(results, [postABC, postAC, postAB])
+
+    results = similar.find_similar_posts(postBC, 2)
+    verify_posts(results, [postABC, postAC])
+
+    results = similar.find_similar_posts(postABC, 10)
+    verify_posts(results, [postBC, postAC, postAB, postA])
+
+    results = similar.find_similar_posts(postA, 10)
+    verify_posts(results, [postABC, postAC, postAB])  # sorted by id
+
+    results = similar.find_similar_posts(postAB, 10)
+    verify_posts(results, [postABC, postBC, postAC, postA])
+
+    results = similar.find_similar_posts(postAC, 10)
+    verify_posts(results, [postABC, postBC, postAB, postA])