[textanalytics] refactor pollers + add TA poller protocol

sdk/textanalytics/azure-ai-textanalytics/CHANGELOG.md

@@ -7,6 +7,8 @@
 - Added `begin_recognize_custom_entities` client method to recognize custom named entities in documents.
 - Added `begin_single_label_classify` client method to perform custom single label classification on documents.
 - Added `begin_multi_label_classify` client method to perform custom multi label classification on documents.
+- Added property `details` on returned poller objects which contain long-running operation metadata.
+- Added `TextAnalysisLROPoller` and `AsyncTextAnalysisLROPoller` protocols to describe the return types from long-running operations.
 
 ### Breaking Changes
 

sdk/textanalytics/azure-ai-textanalytics/azure/ai/textanalytics/__init__.py

@@ -58,7 +58,7 @@
     AnalyzeHealthcareEntitiesAction,
 )
 
-from ._lro import AnalyzeHealthcareEntitiesLROPoller, AnalyzeActionsLROPoller, TextAnalyticsLROPoller
+from ._lro import AnalyzeHealthcareEntitiesLROPoller, AnalyzeActionsLROPoller, TextAnalysisLROPoller
 
 __all__ = [
     "TextAnalyticsApiVersion",
@@ -114,7 +114,7 @@
     "ClassifyDocumentResult",
     "ClassificationCategory",
     "AnalyzeHealthcareEntitiesAction",
-    "TextAnalyticsLROPoller",
+    "TextAnalysisLROPoller",
 ]
 
 __version__ = VERSION

sdk/textanalytics/azure-ai-textanalytics/azure/ai/textanalytics/_lro.py

@@ -6,10 +6,10 @@
 import base64
 import functools
 import json
-import datetime
-from typing import Any, Optional, MutableMapping
 from urllib.parse import urlencode
-from azure.core.polling._poller import PollingReturnType
+from typing import Any, Mapping, Optional, Callable, TypeVar
+from typing_extensions import Protocol, runtime_checkable
+from azure.core.polling import PollingMethod
 from azure.core.exceptions import HttpResponseError
 from azure.core.polling import LROPoller
 from azure.core.polling.base_polling import (
@@ -25,6 +25,43 @@
 _SUCCEEDED = frozenset(["succeeded", "partiallycompleted", "partiallysucceeded"])
 
 
+PollingReturnType = TypeVar("PollingReturnType")
+
+
+@runtime_checkable
+class TextAnalysisLROPoller(Protocol[PollingReturnType]):
+    """Implements a protocol which returned poller objects are consistent with.
+    """
+
+    @property
+    def details(self) -> Mapping[str, Any]:
+        ...
+
+    def polling_method(self) -> PollingMethod[PollingReturnType]:  # pylint: disable=no-self-use
+        ...
+
+    def continuation_token(self) -> str:  # pylint: disable=no-self-use
+        ...
+
+    def status(self) -> str:  # pylint: disable=no-self-use
+        ...
+
+    def result(self, timeout: Optional[int] = None) -> PollingReturnType: # pylint: disable=no-self-use, unused-argument
+        ...
+
+    def wait(self, timeout: Optional[float] = None) -> None:  # pylint: disable=no-self-use, unused-argument
+        ...
+
+    def done(self) -> bool:  # pylint: disable=no-self-use
+        ...
+
+    def add_done_callback(self, func: Callable) -> None:  # pylint: disable=no-self-use, unused-argument
+        ...
+
+    def remove_done_callback(self, func: Callable) -> None:  # pylint: disable=no-self-use, unused-argument
+        ...
+
+
 class TextAnalyticsOperationResourcePolling(OperationResourcePolling):
     def __init__(
         self, operation_location_header="operation-location", show_stats=False
@@ -174,52 +211,26 @@ def polling_method(self) -> AnalyzeHealthcareEntitiesLROPollingMethod:
         return self._polling_method  # type: ignore
 
     @property
-    def created_on(self) -> datetime.datetime:
-        """When your healthcare entities job was created
-
-        :return: When your healthcare entities job was created
-        :rtype: ~datetime.datetime
-        """
-        return self.polling_method().created_on
-
-    @property
-    def expires_on(self) -> datetime.datetime:
-        """When your healthcare entities job will expire
-
-        :return: When your healthcare entities job will expire
-        :rtype: ~datetime.datetime
-        """
-        return self.polling_method().expires_on
-
-    @property
-    def last_modified_on(self) -> datetime.datetime:
-        """When your healthcare entities job was last modified
-
-        :return: When your healthcare entities job was last modified
-        :rtype: ~datetime.datetime
-        """
-        return self.polling_method().last_modified_on
-
-    @property
-    def id(self) -> str:
-        """ID of your call to :func:`begin_analyze_healthcare_entities`
-
-        :return: ID of your call to :func:`begin_analyze_healthcare_entities`
-        :rtype: str
-        """
-        return self.polling_method().id
-
-    @property
-    def display_name(self) -> str:
-        """Given display_name to the healthcare entities job
-
-        :return: Display name of the healthcare entities job.
-        :rtype: str
+    def details(self) -> Mapping[str, Any]:
+        return {
+            "id": self.polling_method().id,
+            "created_on": self.polling_method().created_on,
+            "expires_on": self.polling_method().expires_on,
+            "display_name": self.polling_method().display_name,
+            "last_modified_on": self.polling_method().last_modified_on,
+        }
 
-        .. versionadded:: 2022-04-01-preview
-            *display_name* property.
-        """
-        return self.polling_method().display_name
+    def __getattr__(self, item):
+        attrs = [
+            "created_on",
+            "expires_on",
+            "display_name",
+            "last_modified_on",
+            "id"
+        ]
+        if item in attrs:
+            return self.details[item]
+        return self.__getattribute__(item)
 
     @classmethod
     def from_continuation_token(  # type: ignore
@@ -323,19 +334,19 @@ def display_name(self):
     def actions_failed_count(self):
         if not self._current_body:
             return None
-        return self._current_body.additional_properties["tasks"]["failed"]
+        return self._current_body.additional_properties.get("tasks", {}).get("failed", None)
 
     @property
     def actions_in_progress_count(self):
         if not self._current_body:
             return None
-        return self._current_body.additional_properties["tasks"]["inProgress"]
+        return self._current_body.additional_properties.get("tasks", {}).get("inProgress", None)
 
     @property
     def actions_succeeded_count(self):
         if not self._current_body:
             return None
-        return self._current_body.additional_properties["tasks"]["completed"]
+        return self._current_body.additional_properties.get("tasks", {}).get("completed", None)
 
     @property
     def last_modified_on(self):
@@ -347,7 +358,7 @@ def last_modified_on(self):
     def total_actions_count(self):
         if not self._current_body:
             return None
-        return self._current_body.additional_properties["tasks"]["total"]
+        return self._current_body.additional_properties.get("tasks", {}).get("total", None)
 
     @property
     def id(self):
@@ -370,137 +381,42 @@ def polling_method(self) -> AnalyzeActionsLROPollingMethod:
         return self._polling_method  # type: ignore
 
     @property
-    def created_on(self) -> datetime.datetime:
-        """When your analyze job was created
-
-        :return: When your analyze job was created
-        :rtype: ~datetime.datetime
-        """
-        return self.polling_method().created_on
-
-    @property
-    def expires_on(self) -> datetime.datetime:
-        """When your analyze job will expire
-
-        :return: When your analyze job will expire
-        :rtype: ~datetime.datetime
-        """
-        return self.polling_method().expires_on
-
-    @property
-    def display_name(self) -> Optional[str]:
-        """The display name of your :func:`begin_analyze_actions` call.
-
-        Corresponds to the `display_name` kwarg you pass to your
-        :func:`begin_analyze_actions` call.
-
-        :return: The display name of your :func:`begin_analyze_actions` call.
-        :rtype: str
-        """
-        return self.polling_method().display_name
-
-    @property
-    def actions_failed_count(self) -> int:
-        """Total number of actions that have failed
-
-        :return: Total number of actions that have failed
-        :rtype: int
-        """
-        return self.polling_method().actions_failed_count
-
-    @property
-    def actions_in_progress_count(self) -> int:
-        """Total number of actions currently in progress
-
-        :return: Total number of actions currently in progress
-        :rtype: int
-        """
-        return self.polling_method().actions_in_progress_count
-
-    @property
-    def actions_succeeded_count(self) -> int:
-        """Total number of actions that succeeded
-
-        :return: Total number of actions that succeeded
-        :rtype: int
-        """
-        return self.polling_method().actions_succeeded_count
-
-    @property
-    def last_modified_on(self) -> datetime.datetime:
-        """The last time your actions results were updated
-
-        :return: The last time your actions results were updated
-        :rtype: ~datetime.datetime
-        """
-        return self.polling_method().last_modified_on
-
-    @property
-    def total_actions_count(self) -> int:
-        """Total number of actions you submitted
-
-        :return: Total number of actions submitted
-        :rtype: int
-        """
-        return self.polling_method().total_actions_count
-
-    @property
-    def id(self) -> str:
-        """ID of your :func:`begin_analyze_actions` call.
-
-        :return: ID of your :func:`begin_analyze_actions` call.
-        :rtype: str
-        """
-        return self.polling_method().id
-
-    @classmethod
-    def from_continuation_token(  # type: ignore
-        cls,
-        polling_method: AnalyzeActionsLROPollingMethod,
-        continuation_token: str,
-        **kwargs: Any
-    ) -> "AnalyzeActionsLROPoller":  # type: ignore
-        """
-        :meta private:
-        """
-        client, initial_response, deserialization_callback = polling_method.from_continuation_token(
-            continuation_token, **kwargs
-        )
-        polling_method._lro_algorithms = [  # pylint: disable=protected-access
-            TextAnalyticsOperationResourcePolling(
-                show_stats=initial_response.context.options["show_stats"]
-            )
-        ]
-        return cls(
-            client,
-            initial_response,
-            functools.partial(deserialization_callback, initial_response),
-            polling_method
-        )
-
-
-class TextAnalyticsLROPoller(LROPoller[PollingReturnType]):
-    def polling_method(self) -> AnalyzeActionsLROPollingMethod:
-        """Return the polling method associated to this poller."""
-        return self._polling_method  # type: ignore
-
-    @property
-    def details(self) -> MutableMapping[str, Any]:
+    def details(self) -> Mapping[str, Any]:
         return {
             "id": self.polling_method().id,
             "created_on": self.polling_method().created_on,
             "expires_on": self.polling_method().expires_on,
             "display_name": self.polling_method().display_name,
             "last_modified_on": self.polling_method().last_modified_on,
+            "actions_failed_count": self.polling_method().actions_failed_count,
+            "actions_in_progress_count": self.polling_method().actions_in_progress_count,
+            "actions_succeeded_count": self.polling_method().actions_succeeded_count,
+            "total_actions_count": self.polling_method().total_actions_count,
         }
 
+    def __getattr__(self, item):
+        attrs = [
+            "created_on",
+            "expires_on",
+            "display_name",
+            "actions_failed_count",
+            "actions_in_progress_count",
+            "actions_succeeded_count",
+            "total_actions_count",
+            "last_modified_on",
+            "id"
+        ]
+        if item in attrs:
+            return self.details[item]
+        return self.__getattribute__(item)
+
     @classmethod
     def from_continuation_token(  # type: ignore
         cls,
         polling_method: AnalyzeActionsLROPollingMethod,
         continuation_token: str,
         **kwargs: Any
-    ) -> "TextAnalyticsLROPoller":  # type: ignore
+    ) -> "AnalyzeActionsLROPoller":  # type: ignore
         """
         :meta private:
         """