[formrecognizer] reduce time for recorded tests runs #11970
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -7,6 +7,7 @@ | |
| import six | ||
| from azure.core.credentials import AzureKeyCredential | ||
| from azure.core.pipeline.policies import AzureKeyCredentialPolicy | ||
| from azure.core.pipeline.transport import HttpTransport | ||
| from azure.core.exceptions import ( | ||
| ResourceNotFoundError, | ||
| ResourceExistsError, | ||
|
|
@@ -24,6 +25,30 @@ | |
| } | ||
|
|
||
|
|
||
| class TransportWrapper(HttpTransport): | ||
| """Wrapper class that ensures that an inner client created | ||
| by a `get_client` method does not close the outer transport for the parent | ||
| when used in a context manager. | ||
| """ | ||
| def __init__(self, transport): | ||
| self._transport = transport | ||
|
|
||
| def send(self, request, **kwargs): | ||
| return self._transport.send(request, **kwargs) | ||
|
|
||
| def open(self): | ||
| pass | ||
|
|
||
| def close(self): | ||
| pass | ||
|
|
||
| def __enter__(self): | ||
| pass | ||
|
|
||
| def __exit__(self, *args): # pylint: disable=arguments-differ | ||
| pass | ||
|
|
||
|
|
||
| def get_authentication_policy(credential): | ||
| authentication_policy = None | ||
| if credential is None: | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -71,11 +71,13 @@ def __init__( | |
| ) -> None: | ||
|
|
||
| authentication_policy = get_authentication_policy(credential) | ||
| polling_interval = kwargs.pop("polling_interval", POLLING_INTERVAL) | ||
| self._client = FormRecognizer( | ||
| endpoint=endpoint, | ||
| credential=credential, # type: ignore | ||
| sdk_moniker=USER_AGENT, | ||
| authentication_policy=authentication_policy, | ||
| polling_interval=polling_interval, | ||
| **kwargs | ||
| ) | ||
|
|
||
|
|
@@ -119,7 +121,7 @@ async def begin_recognize_receipts( | |
| :caption: Recognize US sales receipt fields. | ||
| """ | ||
|
|
||
| polling_interval = kwargs.pop("polling_interval", self._client._config.polling_interval) | ||
| continuation_token = kwargs.pop("continuation_token", None) | ||
| content_type = kwargs.pop("content_type", None) | ||
| if content_type == "application/json": | ||
|
|
@@ -176,7 +178,7 @@ async def begin_recognize_receipts_from_url( | |
| :caption: Recognize US sales receipt fields from a URL. | ||
| """ | ||
|
|
||
| polling_interval = kwargs.pop("polling_interval", self._client._config.polling_interval) | ||
| continuation_token = kwargs.pop("continuation_token", None) | ||
| include_text_content = kwargs.pop("include_text_content", False) | ||
|
|
||
|
|
@@ -230,7 +232,7 @@ async def begin_recognize_content( | |
| :caption: Recognize text and content/layout information from a form. | ||
| """ | ||
|
|
||
| polling_interval = kwargs.pop("polling_interval", self._client._config.polling_interval) | ||
| continuation_token = kwargs.pop("continuation_token", None) | ||
| content_type = kwargs.pop("content_type", None) | ||
| if content_type == "application/json": | ||
|
|
@@ -268,7 +270,7 @@ async def begin_recognize_content_from_url(self, form_url: str, **kwargs: Any) - | |
| :raises ~azure.core.exceptions.HttpResponseError: | ||
| """ | ||
|
|
||
| polling_interval = kwargs.pop("polling_interval", self._client._config.polling_interval) | ||
| continuation_token = kwargs.pop("continuation_token", None) | ||
| return await self._client.begin_analyze_layout_async( # type: ignore | ||
| file_stream={"source": form_url}, | ||
|
|
@@ -324,7 +326,7 @@ async def begin_recognize_custom_forms( | |
| raise ValueError("model_id cannot be None or empty.") | ||
|
|
||
| cls = kwargs.pop("cls", None) | ||
| polling_interval = kwargs.pop("polling_interval", self._client._config.polling_interval) | ||
| continuation_token = kwargs.pop("continuation_token", None) | ||
| content_type = kwargs.pop("content_type", None) | ||
| if content_type == "application/json": | ||
|
|
@@ -385,7 +387,7 @@ async def begin_recognize_custom_forms_from_url( | |
| raise ValueError("model_id cannot be None or empty.") | ||
|
|
||
| cls = kwargs.pop("cls", None) | ||
| polling_interval = kwargs.pop("polling_interval", self._client._config.polling_interval) | ||
| continuation_token = kwargs.pop("continuation_token", None) | ||
| include_text_content = kwargs.pop("include_text_content", False) | ||
|
|
||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -15,10 +15,12 @@ | |
| TYPE_CHECKING, | ||
| ) | ||
| from azure.core.polling import AsyncLROPoller | ||
| from azure.core.pipeline import AsyncPipeline | ||
| from azure.core.polling.async_base_polling import AsyncLROBasePolling | ||
| from azure.core.tracing.decorator import distributed_trace | ||
| from azure.core.tracing.decorator_async import distributed_trace_async | ||
| from ._form_recognizer_client_async import FormRecognizerClient | ||
| from ._helpers_async import AsyncTransportWrapper | ||
| from .._generated.aio._form_recognizer_client_async import FormRecognizerClient as FormRecognizer | ||
| from .._generated.models import ( | ||
| TrainRequest, | ||
|
|
@@ -81,13 +83,14 @@ def __init__( | |
| ) -> None: | ||
| self._endpoint = endpoint | ||
| self._credential = credential | ||
| authentication_policy = get_authentication_policy(credential) | ||
| polling_interval = kwargs.pop("polling_interval", POLLING_INTERVAL) | ||
| self._client = FormRecognizer( | ||
| endpoint=self._endpoint, | ||
| credential=self._credential, # type: ignore | ||
| sdk_moniker=USER_AGENT, | ||
| authentication_policy=authentication_policy, | ||
| polling_interval=polling_interval, | ||
| **kwargs | ||
| ) | ||
|
|
||
|
|
@@ -138,7 +141,7 @@ def callback(raw_response): | |
|
|
||
| cls = kwargs.pop("cls", None) | ||
| continuation_token = kwargs.pop("continuation_token", None) | ||
| polling_interval = kwargs.pop("polling_interval", self._client._config.polling_interval) | ||
| deserialization_callback = cls if cls else callback | ||
|
|
||
| if continuation_token: | ||
|
|
@@ -361,7 +364,7 @@ async def begin_copy_model( | |
| raise ValueError("model_id cannot be None or empty.") | ||
|
|
||
| continuation_token = kwargs.pop("continuation_token", None) | ||
| polling_interval = kwargs.pop("polling_interval", self._client._config.polling_interval) | ||
|
|
||
| def _copy_callback(raw_response, _, headers): # pylint: disable=unused-argument | ||
| copy_result = self._client._deserialize(CopyOperationResult, raw_response) | ||
|
|
@@ -395,11 +398,19 @@ def get_form_recognizer_client(self, **kwargs: Any) -> FormRecognizerClient: | |
| :rtype: ~azure.ai.formrecognizer.aio.FormRecognizerClient | ||
| :return: A FormRecognizerClient | ||
| """ | ||
| _pipeline = AsyncPipeline( | ||
| transport=AsyncTransportWrapper(self._client._client._pipeline._transport), | ||
| policies=self._client._client._pipeline._impl_policies | ||
| ) # type: AsyncPipeline | ||
| client = FormRecognizerClient( | ||
| endpoint=self._endpoint, | ||
| credential=self._credential, | ||
| pipeline=_pipeline, | ||
| **kwargs | ||
| ) | ||
| # need to share config, but can't pass as a keyword into client | ||
| client._client._config = self._client._client._config | ||
| return client | ||
|
|
||
| async def __aenter__(self) -> "FormTrainingClient": | ||
| await self._client.__aenter__() | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,31 @@ | ||
| # coding=utf-8 | ||
| # ------------------------------------ | ||
| # Copyright (c) Microsoft Corporation. | ||
| # Licensed under the MIT License. | ||
| # ------------------------------------ | ||
|
|
||
| from azure.core.pipeline.transport import AsyncHttpTransport | ||
|
|
||
|
|
||
| class AsyncTransportWrapper(AsyncHttpTransport): | ||
| """Wrapper class that ensures that an inner client created | ||
| by a `get_client` method does not close the outer transport for the parent | ||
| when used in a context manager. | ||
| """ | ||
| def __init__(self, async_transport): | ||
| self._transport = async_transport | ||
|
|
||
| async def send(self, request, **kwargs): | ||
| return await self._transport.send(request, **kwargs) | ||
|
|
||
| async def open(self): | ||
| pass | ||
|
|
||
| async def close(self): | ||
| pass | ||
|
|
||
| async def __aenter__(self): | ||
| pass | ||
|
|
||
| async def __aexit__(self, *args): # pylint: disable=arguments-differ | ||
| pass |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Are we documenting this
polling_intervalkwarg to the client for users? I think a good idea could be to call it_polling_intervalif we're the only people expecting the use it, but I don't feel too strongly about thisThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is an azure-core setting, so hopefully once Xiang documents all the core keywords we'll be able to link to that :)
Edit: this should be available to users