use ParamSpec to enable intellisense with distributed_trace decorators

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

@@ -27,37 +27,30 @@
 
 import functools
 
-from typing import overload
-
+from typing import Callable, Any, TypeVar, overload
+from typing_extensions import ParamSpec
 from .common import change_context, get_function_and_class_name
 from ..settings import settings
 
-try:
-    from typing import TYPE_CHECKING
-except ImportError:
-    TYPE_CHECKING = False
-
-if TYPE_CHECKING:
-    from typing import Callable, Dict, Optional, Any, TypeVar
 
-    T = TypeVar("T")
+P = ParamSpec("P")
+T = TypeVar("T")
 
 
 @overload
-def distributed_trace(__func):
-    # type: (Callable[..., T]) -> Callable[..., T]
+def distributed_trace(__func: Callable[P, T]) -> Callable[P, T]:
     pass
 
 
 @overload
-def distributed_trace(**kwargs):  # pylint:disable=function-redefined,unused-argument
-    # type: (**Any) -> Callable[[Callable[..., T]], Callable[..., T]]
+def distributed_trace(  # pylint:disable=function-redefined
+    **kwargs: Any,  # pylint:disable=unused-argument
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
     pass
 
 
 def distributed_trace(  # pylint:disable=function-redefined
-    __func=None,  # type: Callable[..., T]
-    **kwargs  # type: Any
+    __func: Callable[P, T] = None, **kwargs: Any
 ):
     """Decorator to apply to function to get traced automatically.
 
@@ -69,12 +62,9 @@ def distributed_trace(  # pylint:disable=function-redefined
     name_of_span = kwargs.pop("name_of_span", None)
     tracing_attributes = kwargs.pop("tracing_attributes", {})
 
-    def decorator(func):
-        # type: (Callable[..., T]) -> Callable[..., T]
-
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
         @functools.wraps(func)
-        def wrapper_use_tracer(*args, **kwargs):
-            # type: (*Any, **Any) -> T
+        def wrapper_use_tracer(*args: Any, **kwargs: Any) -> T:
             merge_span = kwargs.pop("merge_span", False)
             passed_in_parent = kwargs.pop("parent_span", None)
 

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

@@ -27,31 +27,31 @@
 
 import functools
 
-from typing import Awaitable, Callable, Dict, Optional, Any, TypeVar, overload
-
+from typing import Awaitable, Callable, Any, TypeVar, overload
+from typing_extensions import ParamSpec
 from .common import change_context, get_function_and_class_name
 from ..settings import settings
 
-
+P = ParamSpec("P")
 T = TypeVar("T")
 
 
 @overload
 def distributed_trace_async(
-    __func: Callable[..., Awaitable[T]]
-) -> Callable[..., Awaitable[T]]:
+    __func: Callable[P, Awaitable[T]]
+) -> Callable[P, Awaitable[T]]:
     pass
 
 
 @overload
 def distributed_trace_async(  # pylint:disable=function-redefined
-    **kwargs: Any  # pylint:disable=unused-argument
-) -> Callable[[Callable[..., Awaitable[T]]], Callable[..., Awaitable[T]]]:
+    **kwargs: Any,  # pylint:disable=unused-argument
+) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
     pass
 
 
 def distributed_trace_async(  # pylint:disable=function-redefined
-    __func: Callable[..., Awaitable[T]] = None, **kwargs: Any
+    __func: Callable[P, Awaitable[T]] = None, **kwargs: Any
 ):
     """Decorator to apply to function to get traced automatically.
 
@@ -63,7 +63,7 @@ def distributed_trace_async(  # pylint:disable=function-redefined
     name_of_span = kwargs.pop("name_of_span", None)
     tracing_attributes = kwargs.pop("tracing_attributes", {})
 
-    def decorator(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
+    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
         @functools.wraps(func)
         async def wrapper_use_tracer(*args: Any, **kwargs: Any) -> T:
             merge_span = kwargs.pop("merge_span", False)

sdk/core/azure-core/dev_requirements.txt

@@ -1,6 +1,5 @@
 trio
 aiohttp>=3.0
-typing_extensions>=3.7.2
 opencensus>=0.6.0
 opencensus-ext-azure
 opencensus-ext-threading

sdk/core/azure-core/setup.py

@@ -69,5 +69,6 @@
     install_requires=[
         'requests>=2.18.4',
         'six>=1.11.0',
+        "typing-extensions>=4.0.1"
     ],
 )

shared_requirements.txt

@@ -129,6 +129,7 @@ isodate>=0.6.0
 avro>=1.11.0
 pyjwt>=1.7.1
 chardet<5,>=3.0.2
+#override azure-core typing-extensions>=4.0.1
 #override azure-search-documents typing-extensions>=3.7.4.3
 #override azure azure-keyvault~=1.0
 #override azure-mgmt-core azure-core<2.0.0,>=1.15.0