Skip to content

Commit

Permalink
use ParamSpec to enable intellisense with distributed_trace decorators (
Browse files Browse the repository at this point in the history
Azure#22891)

* use ParamSpec in distributed_trace decorators to enable intellisense in py3.10+

* use backported ParamSpec from typings

* add dependency on typing_extensions

* fix name, remove from dev_reqs

* fix typing-extensions version

* update changelog
  • Loading branch information
kristapratico authored and rakshith91 committed Apr 7, 2022
1 parent 8f6897a commit ca3ad67
Show file tree
Hide file tree
Showing 6 changed files with 23 additions and 31 deletions.
1 change: 1 addition & 0 deletions sdk/core/azure-core/CHANGELOG.md
Expand Up @@ -3,6 +3,7 @@
## 1.23.0 (Unreleased)

### Features Added
- Improve intellisense type hinting for service client methods #22891

- Add a case insensitive dict `case_insensitive_dict` in `azure.core.utils`. #23206

Expand Down
32 changes: 11 additions & 21 deletions sdk/core/azure-core/azure/core/tracing/decorator.py
Expand Up @@ -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.
Expand All @@ -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)

Expand Down
18 changes: 9 additions & 9 deletions sdk/core/azure-core/azure/core/tracing/decorator_async.py
Expand Up @@ -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.
Expand All @@ -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)
Expand Down
1 change: 0 additions & 1 deletion 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
Expand Down
1 change: 1 addition & 0 deletions sdk/core/azure-core/setup.py
Expand Up @@ -69,5 +69,6 @@
install_requires=[
'requests>=2.18.4',
'six>=1.11.0',
"typing-extensions>=4.0.1"
],
)
1 change: 1 addition & 0 deletions shared_requirements.txt
Expand Up @@ -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
Expand Down

0 comments on commit ca3ad67

Please sign in to comment.