Merge pull request #1039 from adamtheturtle/query-2

Add support for `max_num_results`

Author: adamtheturtle

src/vws/exceptions.py

@@ -5,6 +5,28 @@
 from requests import Response
 
 
+class MaxNumResultsOutOfRange(Exception):
+    """
+    Exception raised when the ``max_num_results`` given to the Cloud
+    Recognition Web API query endpoint is out of range.
+    """
+
+    def __init__(self, response: Response) -> None:
+        """
+        Args:
+            response: The response to a request to Vuforia.
+        """
+        super().__init__(response.text)
+        self._response = response
+
+    @property
+    def response(self) -> Response:
+        """
+        The response returned by Vuforia which included this error.
+        """
+        return self._response
+
+
 class UnknownTarget(Exception):
     """
     Exception raised when Vuforia returns a response with a result code

src/vws/query.py

@@ -11,6 +11,7 @@
 
 from ._authorization import authorization_header, rfc_1123_date
 from ._result_codes import raise_for_result_code
+from .exceptions import MaxNumResultsOutOfRange
 
 
 class CloudRecoService:
@@ -34,7 +35,11 @@ def __init__(
         self._client_secret_key = client_secret_key.encode()
         self._base_vwq_url = base_vwq_url
 
-    def query(self, image: io.BytesIO) -> List[Dict[str, Any]]:
+    def query(
+        self,
+        image: io.BytesIO,
+        max_num_results: int = 1,
+    ) -> List[Dict[str, Any]]:
         """
         Use the Vuforia Web Query API to make an Image Recognition Query.
 
@@ -51,6 +56,7 @@ def query(self, image: io.BytesIO) -> List[Dict[str, Any]]:
         image_content = image.getvalue()
         body = {
             'image': ('image.jpeg', image_content, 'image/jpeg'),
+            'max_num_results': (None, int(max_num_results), 'text/plain'),
         }
         date = rfc_1123_date()
         request_path = '/v1/query'
@@ -81,6 +87,9 @@ def query(self, image: io.BytesIO) -> List[Dict[str, Any]]:
             data=content,
         )
 
+        if 'Integer out of range' in response.text:
+            raise MaxNumResultsOutOfRange(response=response)
+
         raise_for_result_code(
             response=response,
             expected_result_code='Success',

tests/test_query.py

@@ -5,10 +5,12 @@
 import io
 import uuid
 
+import pytest
 from mock_vws import MockVWS
 from mock_vws.database import VuforiaDatabase
 
 from vws import VWS, CloudRecoService
+from vws.exceptions import MaxNumResultsOutOfRange
 
 
 class TestQuery:
@@ -111,3 +113,57 @@ def test_default(
         vws_client.wait_for_target_processed(target_id=target_id_2)
         matches = cloud_reco_client.query(image=high_quality_image)
         assert len(matches) == 1
+
+    def test_custom(
+        self,
+        vws_client: VWS,
+        cloud_reco_client: CloudRecoService,
+        high_quality_image: io.BytesIO,
+    ) -> None:
+        """
+        It is possible to set a custom ``max_num_results``.
+        """
+        target_id = vws_client.add_target(
+            name=uuid.uuid4().hex,
+            width=1,
+            image=high_quality_image,
+        )
+        target_id_2 = vws_client.add_target(
+            name=uuid.uuid4().hex,
+            width=1,
+            image=high_quality_image,
+        )
+        target_id_3 = vws_client.add_target(
+            name=uuid.uuid4().hex,
+            width=1,
+            image=high_quality_image,
+        )
+        vws_client.wait_for_target_processed(target_id=target_id)
+        vws_client.wait_for_target_processed(target_id=target_id_2)
+        vws_client.wait_for_target_processed(target_id=target_id_3)
+        matches = cloud_reco_client.query(
+            image=high_quality_image,
+            max_num_results=2,
+        )
+        assert len(matches) == 2
+
+    def test_too_many(
+        self,
+        cloud_reco_client: CloudRecoService,
+        high_quality_image: io.BytesIO,
+    ) -> None:
+        """
+        A ``MaxNumResultsOutOfRange`` error is raised if the given
+        ``max_num_results`` is out of range.
+        """
+        with pytest.raises(MaxNumResultsOutOfRange) as exc:
+            cloud_reco_client.query(
+                image=high_quality_image,
+                max_num_results=51,
+            )
+
+        expected_value = (
+            "Integer out of range (51) in form data part 'max_result'. "
+            'Accepted range is from 1 to 50 (inclusive).'
+        )
+        assert str(exc.value) == exc.value.response.text == expected_value