Improve docs search ranking for API queries

docs/_static/custom.js

@@ -329,7 +329,134 @@ function initShibuyaRightToc() {
   syncRightTocCodeButtons(localtoc);
 }
 
+const upltApiSearchGenericTerms = new Set([
+  "api",
+  "apis",
+  "attribute",
+  "attributes",
+  "class",
+  "classes",
+  "doc",
+  "docs",
+  "documentation",
+  "function",
+  "functions",
+  "method",
+  "methods",
+  "object",
+  "objects",
+  "reference",
+  "references",
+]);
+
+function normalizeApiSearchTerm(term) {
+  return String(term || "")
+    .toLowerCase()
+    .replace(/\(\)$/, "")
+    .trim();
+}
+
+function isGenericApiSearchTerm(term) {
+  return upltApiSearchGenericTerms.has(normalizeApiSearchTerm(term));
+}
+
+function getApiSearchTerms(terms) {
+  if (terms instanceof Set) {
+    return Array.from(terms);
+  }
+  return Array.from(terms || []);
+}
+
+function apiSearchResultMatchesQueryTerm(title, anchor, terms) {
+  const haystack = `${title || ""} ${anchor || ""}`.toLowerCase();
+  const leaf = haystack.split("#").pop().split(".").pop();
+  return getApiSearchTerms(terms).some((term) => {
+    const normalized = normalizeApiSearchTerm(term);
+    if (!normalized || isGenericApiSearchTerm(normalized)) return false;
+    return (
+      leaf === normalized ||
+      leaf.includes(normalized) ||
+      haystack.includes("." + normalized)
+    );
+  });
+}
+
+function initApiSearchScoring() {
+  if (typeof Search === "undefined" || typeof Scorer === "undefined") return;
+  if (Search.upltApiSearchScoring === "1") return;
+
+  const previousParseQuery = Search._parseQuery;
+  if (typeof previousParseQuery === "function") {
+    Search._parseQuery = function (query) {
+      const parsed = previousParseQuery.call(this, query);
+      const queryTerms = new Set(
+        getApiSearchTerms(parsed && parsed[4])
+          .map(normalizeApiSearchTerm)
+          .filter(Boolean),
+      );
+      Search.upltQueryTerms = queryTerms;
+      Search.upltApiLikeQuery =
+        /[.()]/.test(query || "") ||
+        getApiSearchTerms(queryTerms).some(isGenericApiSearchTerm);
+      return parsed;
+    };
+  }
+
+  const previousObjectSearch = Search.performObjectSearch;
+  if (typeof previousObjectSearch === "function") {
+    Search.performObjectSearch = function (object, objectTerms) {
+      const normalizedObject = normalizeApiSearchTerm(object);
+      const filteredTerms = new Set(
+        getApiSearchTerms(objectTerms)
+          .map(normalizeApiSearchTerm)
+          .filter((term) => term && !isGenericApiSearchTerm(term)),
+      );
+      if (normalizedObject && !isGenericApiSearchTerm(normalizedObject)) {
+        filteredTerms.add(normalizedObject);
+      }
+      return previousObjectSearch.call(this, object, filteredTerms);
+    };
+  }
+
+  const previousScore = Scorer.score;
+  Scorer.score = function (result) {
+    let score =
+      typeof previousScore === "function" ? previousScore(result) : result[4];
+    if (!Number.isFinite(score)) {
+      score = Number.isFinite(result[4]) ? result[4] : 0;
+    }
+
+    const [docname, title, anchor, descr, _baseScore, _filename, kind] = result;
+    const isApiReference = String(docname || "").startsWith("api/");
+    const isApiLikeQuery = !!Search.upltApiLikeQuery;
+    const queryTerms = Search.upltQueryTerms || new Set();
+
+    if (isApiReference && kind === "object") {
+      score += 24;
+      if (isApiLikeQuery) score += 16;
+      if (apiSearchResultMatchesQueryTerm(title, anchor, queryTerms)) {
+        score += 12;
+      }
+      if (
+        queryTerms.has("function") &&
+        String(descr || "").toLowerCase().includes("python function")
+      ) {
+        score += 4;
+      }
+    } else if (isApiReference && isApiLikeQuery) {
+      score += kind === "title" || kind === "index" ? 12 : 8;
+    } else if (!isApiReference && isApiLikeQuery) {
+      score -= 4;
+    }
+
+    return score;
+  };
+
+  Search.upltApiSearchScoring = "1";
+}
+
 document.addEventListener("DOMContentLoaded", function () {
+  initApiSearchScoring();
   initScrollChromeFade();
 
   if (document.querySelector(".sphx-glr-thumbcontainer")) {

ultraplot/tests/test_docs_search.py

@@ -0,0 +1,119 @@
+import json
+import shutil
+import subprocess
+import textwrap
+from pathlib import Path
+
+import pytest
+
+ROOT = Path(__file__).resolve().parents[2]
+
+
+def test_docs_search_prioritizes_api_references_for_generic_function_queries():
+    if shutil.which("node") is None:
+        pytest.skip("Node.js is required to exercise the docs search JavaScript.")
+
+    script = textwrap.dedent(r"""
+        const fs = require("fs");
+        const vm = require("vm");
+        const listeners = [];
+        const classList = {
+          add() {},
+          remove() {},
+          contains() { return false; },
+          toggle() {},
+        };
+        const context = {
+          console,
+          Scorer: {},
+          Search: {
+            _parseQuery(query) {
+              const objectTerms = new Set(query.toLowerCase().split(/\s+/).filter(Boolean));
+              return [query, new Set(), new Set(), new Set(), objectTerms];
+            },
+            performObjectSearch(_object, objectTerms) {
+              this.lastObjectTerms = Array.from(objectTerms);
+              return this.lastObjectTerms;
+            },
+          },
+        };
+        context.window = {
+          innerWidth: 1024,
+          location: { hash: "", pathname: "/search.html" },
+          requestAnimationFrame(callback) { callback(); },
+          scrollY: 0,
+          addEventListener() {},
+        };
+        context.document = {
+          body: {
+            classList,
+            dataset: {},
+            appendChild() {},
+            getAttribute() { return ""; },
+          },
+          documentElement: { classList, dataset: {} },
+          addEventListener(type, callback) {
+            if (type === "DOMContentLoaded") listeners.push(callback);
+          },
+          querySelector() { return null; },
+          querySelectorAll() { return []; },
+        };
+        context.localStorage = {
+          getItem() { return null; },
+          setItem() {},
+        };
+
+        vm.runInNewContext(fs.readFileSync("docs/_static/custom.js", "utf8"), context);
+        for (const callback of listeners) callback.call(context.document);
+
+        const parsed = context.Search._parseQuery("format function");
+        const filteredTerms = context.Search.performObjectSearch("format", parsed[4]);
+        const apiObjectScore = context.Scorer.score([
+          "api/ultraplot.axes.Axes",
+          "ultraplot.axes.Axes.format",
+          "#ultraplot.axes.Axes.format",
+          "Python method, in Axes",
+          16,
+          "api/ultraplot.axes.Axes.html",
+          "object",
+        ]);
+        const apiTextScore = context.Scorer.score([
+          "api/ultraplot.axes.Axes",
+          "Axes",
+          "",
+          null,
+          16,
+          "api/ultraplot.axes.Axes.html",
+          "text",
+        ]);
+        const guideScore = context.Scorer.score([
+          "basics",
+          "The basics",
+          "",
+          null,
+          16,
+          "basics.html",
+          "text",
+        ]);
+
+        console.log(JSON.stringify({
+          apiLikeQuery: context.Search.upltApiLikeQuery,
+          apiObjectScore,
+          apiTextScore,
+          filteredTerms,
+          guideScore,
+        }));
+        """)
+
+    result = subprocess.run(
+        ["node", "-e", script],
+        cwd=ROOT,
+        text=True,
+        capture_output=True,
+        check=True,
+    )
+    data = json.loads(result.stdout)
+
+    assert data["apiLikeQuery"] is True
+    assert data["filteredTerms"] == ["format"]
+    assert data["apiObjectScore"] > data["apiTextScore"] > data["guideScore"]