From 17c124b9c6ea924c94c19bdb70682f71c5a7fe17 Mon Sep 17 00:00:00 2001
From: Sergei Rybakov <koncopd@gmail.com>
Date: Fri, 22 Nov 2024 16:57:09 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=8E=A8=20Make=20search=20consistent=20wit?=
 =?UTF-8?q?h=20the=20lamindb=20implementation=20(#170)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .gitignore                      |  5 +++++
 bionty/base/_public_ontology.py | 29 +++++++++++++++--------------
 pyproject.toml                  |  1 +
 3 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/.gitignore b/.gitignore
index 291058c..07146bc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,8 @@
+# bionty data
+data/
+sources/
+versions/
+
 # macOS
 .DS_Store
 .AppleDouble
diff --git a/bionty/base/_public_ontology.py b/bionty/base/_public_ontology.py
index 12e433b..2acbbd7 100644
--- a/bionty/base/_public_ontology.py
+++ b/bionty/base/_public_ontology.py
@@ -561,24 +561,18 @@ def search(
         self,
         string: str,
         *,
-        field: PublicOntologyField | str | None = None,
+        field: PublicOntologyField | str | list[PublicOntologyField | str] = None,
         limit: int | None = None,
         case_sensitive: bool = False,
-        synonyms_field: PublicOntologyField | str | None = "synonyms",
     ):
-        """Search a given string against a PublicOntology field.
+        """Search a given string against a PublicOntology field or fields.
 
         Args:
             string: The input string to match against the field values.
-            field: The PublicOntologyField of the ontology the input string is matching against.
-            top_hit: Return all entries ranked by matching ratios.
-                If True, only return the top match.
-                Defaults to False.
-            limit: Maximum amount of top results to return.
-                   If None, return all results.
-                   Defaults to None.
+            field: The PublicOntologyField or several fileds of the ontology
+                the input string is matching against. Search all fields containing strings by default.
+            limit: Maximum amount of top results to return. If None, return all results.
             case_sensitive: Whether the match is case sensitive.
-            synonyms_field: By default also search against the synonyms (If None, skips search).
 
         Returns:
             Ranked search results.
@@ -590,14 +584,21 @@ def search(
         """
         from lamin_utils._search import search
 
-        return search(
+        if isinstance(field, PublicOntologyField):
+            field = field.name
+        elif field is not None and not isinstance(field, str):
+            field = [f.name if isinstance(f, PublicOntologyField) else f for f in field]
+
+        result = search(
             df=self._df,
             string=string,
-            field=self._get_default_field(field),
+            field=field,
             limit=limit,
             case_sensitive=case_sensitive,
-            synonyms_field=str(synonyms_field),
         )
+        if "ontology_id" in result.columns:
+            result = result.set_index("ontology_id")
+        return result
 
     def diff(
         self, compare_to: PublicOntology, **kwargs
diff --git a/pyproject.toml b/pyproject.toml
index ac63464..02ab984 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,6 +18,7 @@ classifiers = [
 dependencies = [
     "lamindb",
     "lamindb_setup>=0.77.6",
+    "lamin_utils>=0.13.9",
     "filelock",
     "requests",
     "pyyaml",