Refactoring of autocomplete

Author: cthoyt

src/bioregistry/app/api.py

@@ -13,7 +13,7 @@
 from fastapi.responses import JSONResponse, Response
 from pydantic import BaseModel, Field
 
-from .utils import FORMAT_MAP, _autocomplete, _search
+from .utils import FORMAT_MAP, AutocompleteResult, PrefixSearchResult, _autocomplete, _search
 from ..export.rdf_export import (
     collection_to_rdf_str,
     metaresource_to_rdf_str,
@@ -721,34 +721,37 @@ def generate_context_json_ld(
     )
 
 
-@api_router.get("/autocomplete", tags=["search"])
+@api_router.get("/autocomplete", tags=["search"], response_model_exclude_none=True)
 def autocomplete(
     manager: DependsManager,
     q: Annotated[str, Query(description="A query for the prefix")],
-) -> JSONResponse:
+) -> AutocompleteResult:
     """Complete a resolution query."""
-    return JSONResponse(_autocomplete(manager, q))
+    return _autocomplete(manager, q)
 
 
-@api_router.get("/search", tags=["search"])
+@api_router.get("/search", tags=["search"], response_model_exclude_none=True)
 def search(
     manager: DependsManager,
     q: Annotated[str, Query(description="A query for the prefix")],
-) -> JSONResponse:
+) -> list[PrefixSearchResult]:
     """Search for a prefix."""
-    return JSONResponse(_search(manager, q))
+    return _search(manager, q)
 
 
 class ReverseURILookupResults(BaseModel):
+    """Results from reverse URI lookup."""
+
     query: str
     resources: list[Resource]
 
 
-@api_router.get("/uri-search", tags=["search"])
+@api_router.get("/reverse", tags=["search"])
 def reverse_uri_lookup(
     manager: DependsManager,
     uri_prefix: Annotated[str, Query(description="A URI prefix for lookup")],
 ) -> ReverseURILookupResults:
+    """Lookup a URI, inspired by https://prefix.cc/reverse."""
     return ReverseURILookupResults(
         query=uri_prefix,
         resources=[

src/bioregistry/app/utils.py

@@ -3,10 +3,12 @@
 from __future__ import annotations
 
 import json
-from collections.abc import Callable, Mapping, Sequence
+from collections import defaultdict
+from collections.abc import Callable, Sequence
 from functools import partial
-from typing import Any, cast
+from typing import cast
 
+import curies
 import werkzeug
 import yaml
 from flask import (
@@ -71,17 +73,48 @@ def _normalize_prefix_or_404(
     return norm_prefix
 
 
-def _search(manager_: Manager, q: str) -> list[tuple[str, str]]:
-    q_norm = _norm(q)
-    results = [
-        (prefix, lookup if _norm(prefix) != lookup else "")
-        for lookup, prefix in manager_.synonyms.items()
-        if q_norm in lookup
+class PrefixLookup(BaseModel):
+    """A prefix lookup."""
+
+    lookup: str
+    position: int
+    complete: bool
+
+
+class PrefixSearchResult(BaseModel):
+    """Results from searching for a prefix."""
+
+    prefix: str
+    lookups: list[PrefixLookup]
+
+
+def _search(manager_: Manager, query: str) -> list[PrefixSearchResult]:
+    query = _norm(query)
+    prefix_to_lookups = defaultdict(list)
+    for lookup, prefix in manager_.synonyms.items():
+        if query in lookup:
+            start = lookup.find(query)
+            prefix_to_lookups[prefix].append(
+                PrefixLookup(lookup=lookup, position=start, complete=query == lookup)
+            )
+    return [
+        PrefixSearchResult(prefix=prefix, lookups=lookups)
+        for prefix, lookups in prefix_to_lookups.items()
     ]
-    return sorted(results)
 
 
-def _autocomplete(manager_: Manager, q: str, base_url: str | None = None) -> Mapping[str, Any]:
+class AutocompleteResult(BaseModel):
+    """Results from autocomplete."""
+
+    query: str
+    success: bool
+    reason: str
+    results: curies.Reference | list[PrefixSearchResult]
+    url: str | None = None
+    pattern: str | None = None
+
+
+def _autocomplete(manager_: Manager, q: str, base_url: str | None = None) -> AutocompleteResult:
     r"""Run the autocomplete algorithm.
 
     :param manager_: A manager
@@ -130,23 +163,23 @@ def _autocomplete(manager_: Manager, q: str, base_url: str | None = None) -> Map
         else:
             reason = "searched prefix"
             url = None
-        return {
-            "query": q,
-            "results": _search(manager_, q),
-            "success": True,
-            "reason": reason,
-            "url": url,
-        }
+        return AutocompleteResult(
+            query=q,
+            success=True,
+            reason=reason,
+            url=url,
+            results=_search(manager_, q),
+        )
+
     prefix, identifier = q.split(":", 1)
     resource = manager_.get_resource(prefix)
     if resource is None:
-        return {
-            "query": q,
-            "prefix": prefix,
-            "identifier": identifier,
-            "success": False,
-            "reason": "bad prefix",
-        }
+        return AutocompleteResult(
+            query=q,
+            success=False,
+            reason="bad prefix",
+            results=curies.Reference(prefix=prefix, identifier=identifier),
+        )
     pattern = manager_.get_pattern(prefix)
     if pattern is None:
         success = True
@@ -162,15 +195,15 @@ def _autocomplete(manager_: Manager, q: str, base_url: str | None = None) -> Map
         success = False
         reason = "failed validation"
         url = None
-    return {
-        "query": q,
-        "prefix": prefix,
-        "pattern": pattern,
-        "identifier": identifier,
-        "success": success,
-        "reason": reason,
-        "url": url,
-    }
+
+    return AutocompleteResult(
+        query=q,
+        success=success,
+        reason=reason,
+        url=url,
+        results=curies.Reference(prefix=prefix, identifier=identifier),
+        pattern=pattern,
+    )
 
 
 def serialize(