1import functools
2from typing import Optional, Callable, TypeVar, TYPE_CHECKING
3
4if TYPE_CHECKING:
5 from mypy_extensions import VarArg, KwArg # noqa
6 from typing import Any # noqa
7
8AnyCallable = TypeVar("AnyCallable", bound=Callable)
9
10
11def _Tracer_wrap(
12 self, # this will be bound to a Tracer instance
13 name: Optional[str] = None,
14 service: Optional[str] = None,
15 resource: Optional[str] = None,
16 span_type: Optional[str] = None,
17):
18 # type: (...) -> Callable[[VarArg(Any), KwArg(Any)], Any]
19 """
20 A function returning a decorator used to trace an entire function. If the traced function
21 is a coroutine, it traces the coroutine execution when is awaited.
22 If a ``wrap_executor`` callable has been provided in the ``Tracer.configure()``
23 method, it will be called instead of the default one when the function
24 decorator is invoked.
25
26 >>> @tracer.wrap("my.wrapped.function", service="my.service")
27 def run():
28 return "run"
29
30 >>> # name will default to "execute" if unset
31 @tracer.wrap()
32 def execute():
33 return "executed"
34
35 >>> # or use it in asyncio coroutines
36 @tracer.wrap()
37 async def coroutine():
38 return "executed"
39
40 >>> @tracer.wrap()
41 @asyncio.coroutine
42 def coroutine():
43 return "executed"
44
45 You can access the current span using `tracer.current_span()` to set
46 tags:
47
48 >>> @tracer.wrap()
49 def execute():
50 span = tracer.current_span()
51 span.set_tag("a", "b")
52 """
53
54 def wrap_decorator(f: AnyCallable):
55 # type: (...) -> Callable[[VarArg(Any), KwArg(Any)], Any]
56 @functools.wraps(f)
57 def func_wrapper(*args, **kwargs):
58 with self.trace(
59 name or "%s.%s" % (f.__module__, f.__name__),
60 service=service,
61 resource=resource,
62 span_type=span_type,
63 ):
64 return f(*args, **kwargs)
65
66 return func_wrapper
67
68 return wrap_decorator