From eefd64c63c8a3c1baa455674474e90c7c8e60664 Mon Sep 17 00:00:00 2001
From: Ondrej Slama <ondrej.slama@rossum.ai>
Date: Thu, 24 Aug 2023 15:53:24 +0200
Subject: [PATCH] feat(elis_api_clients): add new `search_for_annotations`
 method

---
 rossum_api/elis_api_client.py             | 22 +++++++++++++++
 rossum_api/elis_api_client_sync.py        | 15 ++++++++++
 tests/elis_api_client/test_annotations.py | 34 +++++++++++++++++++++++
 3 files changed, 71 insertions(+)

diff --git a/rossum_api/elis_api_client.py b/rossum_api/elis_api_client.py
index 18a3a0c..277283b 100644
--- a/rossum_api/elis_api_client.py
+++ b/rossum_api/elis_api_client.py
@@ -232,6 +232,28 @@ async def list_all_annotations(
         ):
             yield dacite.from_dict(Annotation, a)
 
+    async def search_for_annotations(
+        self,
+        query: Optional[dict] = None,
+        query_string: Optional[dict] = None,
+        ordering: Sequence[str] = (),
+        sideloads: Sequence[str] = (),
+        **kwargs: Any,
+    ) -> AsyncIterable[Annotation]:
+        """https://elis.rossum.ai/api/docs/internal/#search-for-annotations."""
+        if not query and not query_string:
+            raise ValueError("Either query or query_string must be provided")
+        query_json = {}
+        if query:
+            query_json["query"] = query
+        if query_string:
+            query_json["query_string"] = query_string
+
+        async for a in self._http_client.fetch_all(
+            "annotations/search", ordering, sideloads, json=query_json, method="POST", **kwargs
+        ):
+            yield dacite.from_dict(Annotation, a)
+
     async def retrieve_annotation(
         self, annotation_id: int, sideloads: Sequence[str] = ()
     ) -> Annotation:
diff --git a/rossum_api/elis_api_client_sync.py b/rossum_api/elis_api_client_sync.py
index 48be7c3..1e6950a 100644
--- a/rossum_api/elis_api_client_sync.py
+++ b/rossum_api/elis_api_client_sync.py
@@ -230,6 +230,21 @@ def list_all_annotations(
             )
         )
 
+    def search_for_annotations(
+        self,
+        query: Optional[dict] = None,
+        query_string: Optional[dict] = None,
+        ordering: Sequence[str] = (),
+        sideloads: Sequence[str] = (),
+        **kwargs: Any,
+    ) -> Iterable[Annotation]:
+        """https://elis.rossum.ai/api/docs/internal/#search-for-annotations."""
+        return self._iter_over_async(
+            self.elis_api_client.search_for_annotations(
+                query, query_string, ordering, sideloads, **kwargs
+            )
+        )
+
     def retrieve_annotation(self, annotation_id: int, sideloads: Sequence[str] = ()) -> Annotation:
         """https://elis.rossum.ai/api/docs/#retrieve-an-annotation."""
         return self.event_loop.run_until_complete(
diff --git a/tests/elis_api_client/test_annotations.py b/tests/elis_api_client/test_annotations.py
index c71274c..64c2fbe 100644
--- a/tests/elis_api_client/test_annotations.py
+++ b/tests/elis_api_client/test_annotations.py
@@ -178,6 +178,23 @@ async def test_list_all_annotations_with_content_sideloads_without_schema_ids(
 
         assert not http_client.fetch_all.called
 
+    async def test_search_for_annotations(self, elis_client, dummy_annotation, mock_generator):
+        client, http_client = elis_client
+        http_client.fetch_all.return_value = mock_generator(dummy_annotation)
+
+        annotations = client.search_for_annotations({"$and": []}, {"string": "expl"})
+
+        async for a in annotations:
+            assert a == Annotation(**dummy_annotation)
+
+        http_client.fetch_all.assert_called_with(
+            "annotations/search",
+            (),
+            (),
+            json={"query": {"$and": []}, "query_string": {"string": "expl"}},
+            method="POST",
+        )
+
     async def test_retrieve_annotation(self, elis_client, dummy_annotation):
         client, http_client = elis_client
         http_client.fetch_one.return_value = dummy_annotation
@@ -312,6 +329,23 @@ def test_list_all_annotations_with_content_sideloads_without_schema_ids(
 
         assert not http_client.fetch_all.called
 
+    def test_search_for_annotations(self, elis_client_sync, dummy_annotation, mock_generator):
+        client, http_client = elis_client_sync
+        http_client.fetch_all.return_value = mock_generator(dummy_annotation)
+
+        annotations = client.search_for_annotations({"$and": []}, {"string": "expl"})
+
+        for a in annotations:
+            assert a == Annotation(**dummy_annotation)
+
+        http_client.fetch_all.assert_called_with(
+            "annotations/search",
+            (),
+            (),
+            json={"query": {"$and": []}, "query_string": {"string": "expl"}},
+            method="POST",
+        )
+
     def test_retrieve_annotation(self, elis_client_sync, dummy_annotation):
         client, http_client = elis_client_sync
         http_client.fetch_one.return_value = dummy_annotation