Azure-Core: Fix all typings that are not pipeline related (#31494)

* Fix all typing that are not pipeline related

* Don't need those quotes

* Update sdk/core/azure-core/azure/core/tracing/_abstract_span.py

Co-authored-by: Kashif Khan <[email protected]>

* Quotes again

* Clean span file a bit

* Simplify AsyncCredentials

---------

Co-authored-by: Kashif Khan <[email protected]>

Author: lmazuel

sdk/core/azure-core/azure/core/credentials_async.py

@@ -2,13 +2,14 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-from typing import Any, Optional
+from __future__ import annotations
+from typing import Any, Optional, AsyncContextManager
 from typing_extensions import Protocol, runtime_checkable
 from .credentials import AccessToken as _AccessToken
 
 
 @runtime_checkable
-class AsyncTokenCredential(Protocol):
+class AsyncTokenCredential(Protocol, AsyncContextManager["AsyncTokenCredential"]):
     """Protocol for classes able to provide OAuth tokens."""
 
     async def get_token(
@@ -30,9 +31,3 @@ async def get_token(
 
     async def close(self) -> None:
         pass
-
-    async def __aenter__(self):
-        pass
-
-    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
-        pass

sdk/core/azure-core/azure/core/messaging.py

@@ -4,6 +4,7 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
+from __future__ import annotations
 import uuid
 from base64 import b64decode
 from datetime import datetime
@@ -139,7 +140,7 @@ def __repr__(self) -> str:
         )[:1024]
 
     @classmethod
-    def from_dict(cls, event: Dict[str, Any]) -> "CloudEvent":
+    def from_dict(cls, event: Dict[str, Any]) -> CloudEvent[DataType]:
         """Returns the deserialized CloudEvent object when a dict is provided.
 
         :param event: The dict representation of the event which needs to be deserialized.
@@ -214,7 +215,7 @@ def from_dict(cls, event: Dict[str, Any]) -> "CloudEvent":
         return event_obj
 
     @classmethod
-    def from_json(cls, event: Any) -> "CloudEvent":
+    def from_json(cls, event: Any) -> CloudEvent[DataType]:
         """
         Returns the deserialized CloudEvent object when a json payload is provided.
         :param event: The json string that should be converted into a CloudEvent. This can also be

sdk/core/azure-core/azure/core/serialization.py

@@ -6,7 +6,7 @@
 # --------------------------------------------------------------------------
 import base64
 from json import JSONEncoder
-from typing import Union, cast
+from typing import Union, cast, Any
 from datetime import datetime, date, time, timedelta
 from datetime import timezone
 
@@ -18,7 +18,7 @@
 class _Null:
     """To create a Falsy object"""
 
-    def __bool__(self):
+    def __bool__(self) -> bool:
         return False
 
 
@@ -115,7 +115,7 @@ def _datetime_as_isostr(dt: Union[datetime, date, time, timedelta]) -> str:
 class AzureJSONEncoder(JSONEncoder):
     """A JSON encoder that's capable of serializing datetime objects and bytes."""
 
-    def default(self, o):  # pylint: disable=too-many-return-statements
+    def default(self, o: Any) -> Any:  # pylint: disable=too-many-return-statements
         if isinstance(o, (bytes, bytearray)):
             return base64.b64encode(o).decode()
         try:

sdk/core/azure-core/azure/core/tracing/_abstract_span.py

@@ -3,19 +3,13 @@
 # Licensed under the MIT License.
 # ------------------------------------
 """Protocol that defines what functions wrappers of tracing libraries should implement."""
+from __future__ import annotations
 from enum import Enum
 from urllib.parse import urlparse
 
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Sequence,
-    Optional,
-    Union,
-    Callable,
-    ContextManager,
-    Dict,
-)
+from typing import Any, Sequence, Optional, Union, Callable, Dict, Type
+from types import TracebackType
+from typing_extensions import Protocol, ContextManager
 from azure.core.pipeline.transport import HttpRequest, HttpResponse, AsyncHttpResponse
 from azure.core.rest import (
     HttpResponse as RestHttpResponse,
@@ -26,23 +20,17 @@
 HttpResponseType = Union[HttpResponse, AsyncHttpResponse, RestHttpResponse, AsyncRestHttpResponse]
 HttpRequestType = Union[HttpRequest, RestHttpRequest]
 
-if TYPE_CHECKING:
-    AttributeValue = Union[
-        str,
-        bool,
-        int,
-        float,
-        Sequence[str],
-        Sequence[bool],
-        Sequence[int],
-        Sequence[float],
-    ]
-    Attributes = Optional[Dict[str, AttributeValue]]
-
-try:
-    from typing_extensions import Protocol
-except ImportError:
-    Protocol = object  # type: ignore
+AttributeValue = Union[
+    str,
+    bool,
+    int,
+    float,
+    Sequence[str],
+    Sequence[bool],
+    Sequence[int],
+    Sequence[float],
+]
+Attributes = Dict[str, AttributeValue]
 
 
 class SpanKind(Enum):
@@ -67,11 +55,11 @@ class AbstractSpan(Protocol):
     """
 
     def __init__(  # pylint: disable=super-init-not-called
-        self, span: Optional[Any] = None, name: Optional[str] = None, **kwargs
+        self, span: Optional[Any] = None, name: Optional[str] = None, **kwargs: Any
     ) -> None:
         pass
 
-    def span(self, name: str = "child_span", **kwargs) -> "AbstractSpan":
+    def span(self, name: str = "child_span", **kwargs: Any) -> AbstractSpan:
         """
         Create a child span for the current span and append it to the child spans list.
         The child span must be wrapped by an implementation of AbstractSpan
@@ -98,10 +86,15 @@ def kind(self, value: SpanKind) -> None:
         :type value: SpanKind
         """
 
-    def __enter__(self):
+    def __enter__(self) -> AbstractSpan:
         """Start a span."""
 
-    def __exit__(self, exception_type, exception_value, traceback):
+    def __exit__(
+        self,
+        exception_type: Optional[Type[BaseException]],
+        exception_value: Optional[BaseException],
+        traceback: TracebackType,
+    ) -> None:
         """Finish a span.
 
         :param exception_type: The type of the exception
@@ -159,7 +152,7 @@ def span_instance(self) -> Any:
         """
 
     @classmethod
-    def link(cls, traceparent: str, attributes: Optional["Attributes"] = None) -> None:
+    def link(cls, traceparent: str, attributes: Optional[Attributes] = None) -> None:
         """
         Given a traceparent, extracts the context and links the context to the current tracer.
 
@@ -170,7 +163,7 @@ def link(cls, traceparent: str, attributes: Optional["Attributes"] = None) -> No
         """
 
     @classmethod
-    def link_from_headers(cls, headers: Dict[str, str], attributes: Optional["Attributes"] = None) -> None:
+    def link_from_headers(cls, headers: Dict[str, str], attributes: Optional[Attributes] = None) -> None:
         """
         Given a dictionary, extracts the context and links the context to the current tracer.
 
@@ -215,7 +208,7 @@ def set_current_tracer(cls, tracer: Any) -> None:
         """
 
     @classmethod
-    def change_context(cls, span: "AbstractSpan") -> ContextManager:
+    def change_context(cls, span: AbstractSpan) -> ContextManager[AbstractSpan]:
         """Change the context for the life of this context manager.
 
         :param span: The span to run in the new context
@@ -235,14 +228,7 @@ def with_current_context(cls, func: Callable) -> Callable:
         """
 
 
-# https://github.com/python/mypy/issues/5837
-if TYPE_CHECKING:
-    _MIXIN_BASE = AbstractSpan
-else:
-    _MIXIN_BASE = object
-
-
-class HttpSpanMixin(_MIXIN_BASE):
+class HttpSpanMixin:
     """Can be used to get HTTP span attributes settings for free."""
 
     _SPAN_COMPONENT = "component"
@@ -253,7 +239,9 @@ class HttpSpanMixin(_MIXIN_BASE):
     _NET_PEER_NAME = "net.peer.name"
     _NET_PEER_PORT = "net.peer.port"
 
-    def set_http_attributes(self, request: HttpRequestType, response: Optional[HttpResponseType] = None) -> None:
+    def set_http_attributes(
+        self: AbstractSpan, request: HttpRequestType, response: Optional[HttpResponseType] = None
+    ) -> None:
         """
         Add correct attributes for a http client span.
 
@@ -262,24 +250,25 @@ def set_http_attributes(self, request: HttpRequestType, response: Optional[HttpR
         :param response: The response received from the server. Is None if no response received.
         :type response: ~azure.core.pipeline.transport.HttpResponse or ~azure.core.pipeline.transport.AsyncHttpResponse
         """
+        # Also see https://github.com/python/mypy/issues/5837
         self.kind = SpanKind.CLIENT
-        self.add_attribute(self._SPAN_COMPONENT, "http")
-        self.add_attribute(self._HTTP_METHOD, request.method)
-        self.add_attribute(self._HTTP_URL, request.url)
+        self.add_attribute(HttpSpanMixin._SPAN_COMPONENT, "http")
+        self.add_attribute(HttpSpanMixin._HTTP_METHOD, request.method)
+        self.add_attribute(HttpSpanMixin._HTTP_URL, request.url)
 
         parsed_url = urlparse(request.url)
         if parsed_url.hostname:
-            self.add_attribute(self._NET_PEER_NAME, parsed_url.hostname)
+            self.add_attribute(HttpSpanMixin._NET_PEER_NAME, parsed_url.hostname)
         if parsed_url.port and parsed_url.port not in [80, 443]:
-            self.add_attribute(self._NET_PEER_PORT, parsed_url.port)
+            self.add_attribute(HttpSpanMixin._NET_PEER_PORT, parsed_url.port)
 
         user_agent = request.headers.get("User-Agent")
         if user_agent:
-            self.add_attribute(self._HTTP_USER_AGENT, user_agent)
+            self.add_attribute(HttpSpanMixin._HTTP_USER_AGENT, user_agent)
         if response and response.status_code:
-            self.add_attribute(self._HTTP_STATUS_CODE, response.status_code)
+            self.add_attribute(HttpSpanMixin._HTTP_STATUS_CODE, response.status_code)
         else:
-            self.add_attribute(self._HTTP_STATUS_CODE, 504)
+            self.add_attribute(HttpSpanMixin._HTTP_STATUS_CODE, 504)
 
 
 class Link:
@@ -291,6 +280,6 @@ class Link:
     :type attributes: dict
     """
 
-    def __init__(self, headers: Dict[str, str], attributes: Optional["Attributes"] = None) -> None:
+    def __init__(self, headers: Dict[str, str], attributes: Optional[Attributes] = None) -> None:
         self.headers = headers
         self.attributes = attributes

sdk/core/azure-core/azure/core/tracing/common.py

@@ -38,7 +38,7 @@
 ]
 
 
-def get_function_and_class_name(func: Callable, *args) -> str:
+def get_function_and_class_name(func: Callable, *args: object) -> str:
     """
     Given a function and its unamed arguments, returns class_name.function_name. It assumes the first argument
     is `self`. If there are no arguments then it only returns the function name.

sdk/core/azure-core/azure/core/tracing/decorator.py

@@ -50,7 +50,9 @@ def distributed_trace(  # pylint:disable=function-redefined
     pass
 
 
-def distributed_trace(__func: Optional[Callable[P, T]] = None, **kwargs: Any):  # pylint:disable=function-redefined
+def distributed_trace(
+    __func: Optional[Callable[P, T]] = None, **kwargs: Any
+) -> Any:  # pylint:disable=function-redefined
     """Decorator to apply to function to get traced automatically.
 
     Span will use the func name or "name_of_span".

sdk/core/azure-core/azure/core/tracing/decorator_async.py

@@ -51,7 +51,7 @@ def distributed_trace_async(  # pylint:disable=function-redefined
 
 def distributed_trace_async(  # pylint:disable=function-redefined
     __func: Optional[Callable[P, Awaitable[T]]] = None, **kwargs: Any
-):
+) -> Any:
     """Decorator to apply to function to get traced automatically.
 
     Span will use the func name or "name_of_span".

sdk/core/azure-core/azure/core/utils/_utils.py

@@ -148,7 +148,7 @@ def __iter__(self) -> Iterator[str]:
     def __len__(self) -> int:
         return len(self._store)
 
-    def lowerkey_items(self):
+    def lowerkey_items(self) -> Iterator[Tuple[str, Any]]:
         return ((lower_case_key, pair[1]) for lower_case_key, pair in self._store.items())
 
     def __eq__(self, other: Any) -> bool: