lamindb 1.4.0__py3-none-any.whl → 1.5.1__py3-none-any.whl

lamindb/models/query_manager.py

@@ -1,18 +1,225 @@
 from __future__ import annotations
 
+import re
+from functools import reduce
 from typing import TYPE_CHECKING, NamedTuple
 
-from django.db import models
+from django.db.models import (
+    IntegerField,
+    Manager,
+    Q,
+    QuerySet,
+    TextField,
+    Value,
+)
+from django.db.models.functions import Cast, Coalesce
+from django.db.models.lookups import (
+    Contains,
+    Exact,
+    IContains,
+    IExact,
+    IRegex,
+    IStartsWith,
+    Regex,
+    StartsWith,
+)
 from lamin_utils import logger
+from lamin_utils._lookup import Lookup
 from lamindb_setup.core._docs import doc_args
 
-from lamindb.models import Record
-
 if TYPE_CHECKING:
-    from lamindb.base.types import StrField
+    from ..base.types import StrField
+
+
+def _search(
+    cls,
+    string: str,
+    *,
+    field: StrField | list[StrField] | None = None,
+    limit: int | None = 20,
+    case_sensitive: bool = False,
+    truncate_string: bool = False,
+) -> QuerySet:
+    """Search.
+
+    Args:
+        string: The input string to match against the field ontology values.
+        field: The field or fields to search. Search all string fields by default.
+        limit: Maximum amount of top results to return.
+        case_sensitive: Whether the match is case sensitive.
+
+    Returns:
+        A sorted `DataFrame` of search results with a score in column `score`.
+        If `return_queryset` is `True`.  `QuerySet`.
+
+    See Also:
+        :meth:`~lamindb.models.Record.filter`
+        :meth:`~lamindb.models.Record.lookup`
+
+    Examples:
+        >>> ulabels = ln.ULabel.from_values(["ULabel1", "ULabel2", "ULabel3"], field="name")
+        >>> ln.save(ulabels)
+        >>> ln.ULabel.search("ULabel2")
+    """
+    if string is None:
+        raise ValueError("Cannot search for None value! Please pass a valid string.")
+
+    input_queryset = (
+        cls.all() if isinstance(cls, (QuerySet, Manager)) else cls.objects.all()
+    )
+    registry = input_queryset.model
+    name_field = getattr(registry, "_name_field", "name")
+    if field is None:
+        fields = [
+            field.name
+            for field in registry._meta.fields
+            if field.get_internal_type() in {"CharField", "TextField"}
+        ]
+    else:
+        if not isinstance(field, list):
+            fields_input = [field]
+        else:
+            fields_input = field
+        fields = []
+        for field in fields_input:
+            if not isinstance(field, str):
+                try:
+                    fields.append(field.field.name)
+                except AttributeError as error:
+                    raise TypeError(
+                        "Please pass a Record string field, e.g., `CellType.name`!"
+                    ) from error
+            else:
+                fields.append(field)
+
+    if truncate_string:
+        if (len_string := len(string)) > 5:
+            n_80_pct = int(len_string * 0.8)
+            string = string[:n_80_pct]
+
+    string = string.strip()
+    string_escape = re.escape(string)
+
+    exact_lookup = Exact if case_sensitive else IExact
+    regex_lookup = Regex if case_sensitive else IRegex
+    contains_lookup = Contains if case_sensitive else IContains
+
+    ranks = []
+    contains_filters = []
+    for field in fields:
+        field_expr = Coalesce(
+            Cast(field, output_field=TextField()),
+            Value(""),
+            output_field=TextField(),
+        )
+        # exact rank
+        exact_expr = exact_lookup(field_expr, string)
+        exact_rank = Cast(exact_expr, output_field=IntegerField()) * 200
+        ranks.append(exact_rank)
+        # exact synonym
+        synonym_expr = regex_lookup(field_expr, rf"(?:^|.*\|){string_escape}(?:\|.*|$)")
+        synonym_rank = Cast(synonym_expr, output_field=IntegerField()) * 200
+        ranks.append(synonym_rank)
+        # match as sub-phrase
+        sub_expr = regex_lookup(
+            field_expr, rf"(?:^|.*[ \|\.,;:]){string_escape}(?:[ \|\.,;:].*|$)"
+        )
+        sub_rank = Cast(sub_expr, output_field=IntegerField()) * 10
+        ranks.append(sub_rank)
+        # startswith and avoid matching string with " " on the right
+        # mostly for truncated
+        startswith_expr = regex_lookup(
+            field_expr, rf"(?:^|.*\|){string_escape}[^ ]*(?:\|.*|$)"
+        )
+        startswith_rank = Cast(startswith_expr, output_field=IntegerField()) * 8
+        ranks.append(startswith_rank)
+        # match as sub-phrase from the left, mostly for truncated
+        right_expr = regex_lookup(field_expr, rf"(?:^|.*[ \|]){string_escape}.*")
+        right_rank = Cast(right_expr, output_field=IntegerField()) * 2
+        ranks.append(right_rank)
+        # match as sub-phrase from the right
+        left_expr = regex_lookup(field_expr, rf".*{string_escape}(?:$|[ \|\.,;:].*)")
+        left_rank = Cast(left_expr, output_field=IntegerField()) * 2
+        ranks.append(left_rank)
+        # simple contains filter
+        contains_expr = contains_lookup(field_expr, string)
+        contains_filter = Q(contains_expr)
+        contains_filters.append(contains_filter)
+        # also rank by contains
+        contains_rank = Cast(contains_expr, output_field=IntegerField())
+        ranks.append(contains_rank)
+        # additional rule for truncated strings
+        # weight matches from the beginning of the string higher
+        # sometimes whole words get truncated and startswith_expr is not enough
+        if truncate_string and field == name_field:
+            startswith_lookup = StartsWith if case_sensitive else IStartsWith
+            name_startswith_expr = startswith_lookup(field_expr, string)
+            name_startswith_rank = (
+                Cast(name_startswith_expr, output_field=IntegerField()) * 2
+            )
+            ranks.append(name_startswith_rank)
 
+    ranked_queryset = (
+        input_queryset.filter(reduce(lambda a, b: a | b, contains_filters))
+        .alias(rank=sum(ranks))
+        .order_by("-rank")
+    )
 
-class QueryManager(models.Manager):
+    return ranked_queryset[:limit]
+
+
+def _lookup(
+    cls,
+    field: StrField | None = None,
+    return_field: StrField | None = None,
+    using_key: str | None = None,
+) -> NamedTuple:
+    """Return an auto-complete object for a field.
+
+    Args:
+        field: The field to look up the values for. Defaults to first string field.
+        return_field: The field to return. If `None`, returns the whole record.
+
+    Returns:
+        A `NamedTuple` of lookup information of the field values with a
+        dictionary converter.
+
+    See Also:
+        :meth:`~lamindb.models.Record.search`
+
+    Examples:
+        >>> import bionty as bt
+        >>> bt.settings.organism = "human"
+        >>> bt.Gene.from_source(symbol="ADGB-DT").save()
+        >>> lookup = bt.Gene.lookup()
+        >>> lookup.adgb_dt
+        >>> lookup_dict = lookup.dict()
+        >>> lookup_dict['ADGB-DT']
+        >>> lookup_by_ensembl_id = bt.Gene.lookup(field="ensembl_gene_id")
+        >>> genes.ensg00000002745
+        >>> lookup_return_symbols = bt.Gene.lookup(field="ensembl_gene_id", return_field="symbol")
+    """
+    from .record import get_name_field
+
+    queryset = cls.all() if isinstance(cls, (QuerySet, Manager)) else cls.objects.all()
+    field = get_name_field(registry=queryset.model, field=field)
+
+    return Lookup(
+        records=queryset,
+        values=[i.get(field) for i in queryset.values()],
+        tuple_name=cls.__class__.__name__,
+        prefix="ln",
+    ).lookup(
+        return_field=(
+            get_name_field(registry=queryset.model, field=return_field)
+            if return_field is not None
+            else None
+        )
+    )
+
+
+# this is the default (._default_manager and ._base_manager) for lamindb models
+class QueryManager(Manager):
     """Manage queries through fields.
 
     See Also:
@@ -64,10 +271,10 @@ class QueryManager(models.Manager):
             >>> label.parents.list("name")
             ['ULabel1', 'ULabel2', 'ULabel3']
         """
-        self._track_run_input_manager()
         if field is None:
             return list(self.all())
         else:
+            self._track_run_input_manager()
             return list(self.values_list(field, flat=True))
 
     def df(self, **kwargs):
@@ -83,29 +290,21 @@ class QueryManager(models.Manager):
         For `**kwargs`, see :meth:`lamindb.models.QuerySet.df`.
         """
         self._track_run_input_manager()
-        return self._all_base_class()
+        return super().all()
 
-    @doc_args(Record.search.__doc__)
+    @doc_args(_search.__doc__)
     def search(self, string: str, **kwargs):
         """{}"""  # noqa: D415
-        from .record import _search
-
         return _search(cls=self.all(), string=string, **kwargs)
 
-    @doc_args(Record.lookup.__doc__)
+    @doc_args(_lookup.__doc__)
     def lookup(self, field: StrField | None = None, **kwargs) -> NamedTuple:
         """{}"""  # noqa: D415
-        from .record import _lookup
-
         return _lookup(cls=self.all(), field=field, **kwargs)
 
+    def get_queryset(self):
+        from .query_set import BasicQuerySet
 
-models.Manager.list = QueryManager.list
-models.Manager.df = QueryManager.df
-models.Manager.search = QueryManager.search
-models.Manager.lookup = QueryManager.lookup
-models.Manager._track_run_input_manager = QueryManager._track_run_input_manager
-# the two lines below would be easy if we could actually inherit; like this,
-# they're suboptimal
-models.Manager._all_base_class = models.Manager.all
-models.Manager.all = QueryManager.all
+        # QueryManager returns BasicQuerySet because it is problematic to redefine .filter and .get
+        # for a query set used by the default manager
+        return BasicQuerySet(model=self.model, using=self._db, hints=self._hints)