diff --git a/vision/google/cloud/vision/decorators.py b/vision/google/cloud/vision/decorators.py index a29e8162afbb..3f44664de231 100644 --- a/vision/google/cloud/vision/decorators.py +++ b/vision/google/cloud/vision/decorators.py @@ -85,6 +85,8 @@ def _create_single_feature_method(feature, enum): image (:class:`~.{module}.types.Image`): The image to analyze. options (:class:`google.gax.CallOptions`): Overrides the default settings for this call, e.g, timeout, retries, etc. + kwargs (dict): Additional properties to be set on the + :class:`~.{module}.types.AnnotateImageRequest`. Returns: :class:`~.{module}.types.AnnotateImageResponse`: The API response. @@ -94,16 +96,17 @@ def _create_single_feature_method(feature, enum): feature_value = {'type': enum.__dict__[feature]} # Define the function to be returned. - def inner(self, image, options=None): + def inner(self, image, options=None, **kwargs): """Return a single feature annotation for the given image. Intended for use with functools.partial, to create the particular single-feature methods. """ - request = { - 'image': image, - 'features': [feature_value], - } + request = dict( + image=image, + features=[feature_value], + **kwargs + ) return self.annotate_image(request, options=options) # Set the appropriate function metadata. diff --git a/vision/tests/unit/test_decorators.py b/vision/tests/unit/test_decorators.py index 8ef86b71ec61..f0841e8ecd21 100644 --- a/vision/tests/unit/test_decorators.py +++ b/vision/tests/unit/test_decorators.py @@ -55,15 +55,21 @@ class SingleFeatureMethodTests(unittest.TestCase): def test_runs_generic_single_image(self, ai): ai.return_value = vision.types.AnnotateImageResponse() + # Prove that other aspects of the AnnotateImageRequest, such as the + # image context, will be preserved. + SENTINEL = object() + # Make a face detection request. client = vision.ImageAnnotatorClient( credentials=mock.Mock(spec=Credentials), ) image = {'source': {'image_uri': 'gs://my-test-bucket/image.jpg'}} - response = client.face_detection(image) + response = client.face_detection(image, image_context=SENTINEL) + assert isinstance(response, vision.types.AnnotateImageResponse) # Assert that the single-image method was called as expected. ai.assert_called_once_with({ 'features': [{'type': vision.enums.Feature.Type.FACE_DETECTION}], 'image': image, + 'image_context': SENTINEL, }, options=None)