@@ -173,6 +173,7 @@ def set_fates(clotho, lachesis, atropos, father="Zues", mother="Themis"):
173173 Any ,
174174 Callable ,
175175 Collection ,
176+ ContextManager ,
176177 Dict ,
177178 Generator ,
178179 Iterable ,
@@ -823,75 +824,110 @@ def extract_text_map(carrier: Dict[str, str]) -> Optional["opentracing.SpanConte
823824# Tracing decorators
824825
825826
826- def trace_with_opname (opname : str ) -> Callable [[Callable [P , R ]], Callable [P , R ]]:
827+ def _create_decorator (
828+ func : Callable [P , R ],
829+ # TODO: What is the correct type for these `Any`? `P.args, P.kwargs` isn't allowed here
830+ wrapping_logic : Callable [[Callable [P , R ], Any , Any ], ContextManager [None ]],
831+ ) -> Callable [P , R ]:
827832 """
828- Decorator to trace a function with a custom opname.
829-
830- See the module's doc string for usage examples.
831-
833+ Creates a decorator that is able to handle sync functions, async functions
834+ (coroutines), and inlineDeferred from Twisted.
835+ Example usage:
836+ ```py
837+ # Decorator to time the function and log it out
838+ def duration(func: Callable[P, R]) -> Callable[P, R]:
839+ @contextlib.contextmanager
840+ def _wrapping_logic(func: Callable[P, R], *args: P.args, **kwargs: P.kwargs):
841+ start_ts = time.time()
842+ yield
843+ end_ts = time.time()
844+ duration = end_ts - start_ts
845+ logger.info("%s took %s seconds", func.__name__, duration)
846+ return _create_decorator(func, _wrapping_logic)
847+ ```
848+ Args:
849+ func: The function to be decorated
850+ wrapping_logic: The business logic of your custom decorator.
851+ This should be a ContextManager so you are able to run your logic
852+ before/after the function as desired.
832853 """
833854
834- def decorator (func : Callable [P , R ]) -> Callable [P , R ]:
835- if opentracing is None :
836- return func # type: ignore[unreachable]
837-
855+ @wraps (func )
856+ async def _wrapper (* args : P .args , ** kwargs : P .kwargs ) -> R :
838857 if inspect .iscoroutinefunction (func ):
839-
840- @wraps (func )
841- async def _trace_inner (* args : P .args , ** kwargs : P .kwargs ) -> R :
842- with start_active_span (opname ):
843- return await func (* args , ** kwargs ) # type: ignore[misc]
844-
858+ with wrapping_logic (func , * args , ** kwargs ):
859+ return await func (* args , ** kwargs )
845860 else :
846861 # The other case here handles both sync functions and those
847862 # decorated with inlineDeferred.
848- @wraps (func )
849- def _trace_inner (* args : P .args , ** kwargs : P .kwargs ) -> R :
850- scope = start_active_span (opname )
851- scope .__enter__ ()
852-
853- try :
854- result = func (* args , ** kwargs )
855- if isinstance (result , defer .Deferred ):
856-
857- def call_back (result : R ) -> R :
858- scope .__exit__ (None , None , None )
859- return result
860-
861- def err_back (result : R ) -> R :
862- scope .__exit__ (None , None , None )
863- return result
864-
865- result .addCallbacks (call_back , err_back )
866-
867- else :
868- if inspect .isawaitable (result ):
869- logger .error (
870- "@trace may not have wrapped %s correctly! "
871- "The function is not async but returned a %s." ,
872- func .__qualname__ ,
873- type (result ).__name__ ,
874- )
863+ scope = wrapping_logic (func , * args , ** kwargs )
864+ scope .__enter__ ()
875865
866+ try :
867+ result = func (* args , ** kwargs )
868+ if isinstance (result , defer .Deferred ):
869+
870+ def call_back (result : R ) -> R :
876871 scope .__exit__ (None , None , None )
872+ return result
877873
878- return result
874+ def err_back (result : R ) -> R :
875+ scope .__exit__ (None , None , None )
876+ return result
879877
880- except Exception as e :
881- scope .__exit__ (type (e ), None , e .__traceback__ )
882- raise
878+ result .addCallbacks (call_back , err_back )
883879
884- return _trace_inner # type: ignore[return-value]
880+ else :
881+ if inspect .isawaitable (result ):
882+ logger .error (
883+ "@trace may not have wrapped %s correctly! "
884+ "The function is not async but returned a %s." ,
885+ func .__qualname__ ,
886+ type (result ).__name__ ,
887+ )
885888
886- return decorator
889+ scope .__exit__ (None , None , None )
890+
891+ return result
892+
893+ except Exception as e :
894+ scope .__exit__ (type (e ), None , e .__traceback__ )
895+ raise
896+
897+ return _wrapper # type: ignore[return-value]
898+
899+
900+ def trace_with_opname (opname : str ) -> Callable [[Callable [P , R ]], Callable [P , R ]]:
901+ """
902+ Decorator to trace a function with a custom opname.
903+ See the module's doc string for usage examples.
904+ """
905+
906+ @contextlib .contextmanager
907+ def _wrapping_logic (func : Callable [P , R ], * args : P .args , ** kwargs : P .kwargs ):
908+ if opentracing is None :
909+ return None
910+
911+ scope = start_active_span (opname )
912+ scope .__enter__ ()
913+ try :
914+ yield
915+ except Exception as e :
916+ scope .__exit__ (type (e ), None , e .__traceback__ )
917+ raise
918+ finally :
919+ scope .__exit__ (None , None , None )
920+
921+ def _decorator (func : Callable [P , R ]):
922+ return _create_decorator (func , _wrapping_logic )
923+
924+ return _decorator
887925
888926
889927def trace (func : Callable [P , R ]) -> Callable [P , R ]:
890928 """
891929 Decorator to trace a function.
892-
893930 Sets the operation name to that of the function's name.
894-
895931 See the module's doc string for usage examples.
896932 """
897933
@@ -900,22 +936,22 @@ def trace(func: Callable[P, R]) -> Callable[P, R]:
900936
901937def tag_args (func : Callable [P , R ]) -> Callable [P , R ]:
902938 """
903- Tags all of the args to the active span.
939+ Decorator to tag all of the args to the active span.
904940 """
905941
906942 if not opentracing :
907943 return func
908944
909- @wraps ( func )
910- def _tag_args_inner ( * args : P .args , ** kwargs : P .kwargs ) -> R :
945+ @contextlib . contextmanager
946+ def _wrapping_logic ( func : Callable [ P , R ], * args : P .args , ** kwargs : P .kwargs ):
911947 argspec = inspect .getfullargspec (func )
912948 for i , arg in enumerate (argspec .args [1 :]):
913949 set_tag ("ARG_" + arg , str (args [i ])) # type: ignore[index]
914950 set_tag ("args" , str (args [len (argspec .args ) :])) # type: ignore[index]
915951 set_tag ("kwargs" , str (kwargs ))
916- return func ( * args , ** kwargs )
952+ yield
917953
918- return _tag_args_inner
954+ return _create_decorator ( func , _wrapping_logic )
919955
920956
921957@contextlib .contextmanager
0 commit comments