Use wrapt instead of functools.wraps in order to preserve function si…

beeline/aiotrace.py

@@ -8,6 +8,8 @@
 import functools
 import inspect
 
+import wrapt
+
 from beeline.trace import Tracer
 
 current_trace_var = contextvars.ContextVar("current_trace")
@@ -72,27 +74,24 @@ def traced_impl(tracer_fn, name, trace_id, parent_id):
     """
     def wrapped(fn):
         if asyncio.iscoroutinefunction(fn):
-            @functools.wraps(fn)
-            async def async_inner(*args, **kwargs):
+            @wrapt.decorator
+            async def async_inner(fn, instance, args, kwargs):
                 with tracer_fn(name=name, trace_id=trace_id, parent_id=parent_id):
                     return await fn(*args, **kwargs)
-
-            return async_inner
+            return async_inner(fn)
         elif inspect.isgeneratorfunction(fn):
-            @functools.wraps(fn)
-            def inner(*args, **kwargs):
+            @wrapt.decorator
+            def inner(fn, instance, args, kwargs):
                 inner_generator = fn(*args, **kwargs)
                 with tracer_fn(name=name, trace_id=trace_id, parent_id=parent_id):
                     yield from inner_generator
-
-            return inner
+            return inner(fn)
         else:
-            @functools.wraps(fn)
-            def inner(*args, **kwargs):
+            @wrapt.decorator
+            def inner(fn, instance, args, kwargs):
                 with tracer_fn(name=name, trace_id=trace_id, parent_id=parent_id):
                     return fn(*args, **kwargs)
-
-            return inner
+            return inner(fn)
 
     return wrapped
 
@@ -108,10 +107,10 @@ def untraced(fn):
 
     # Both synchronous and asynchronous functions may create tasks.
     if asyncio.iscoroutinefunction(fn):
-        @functools.wraps(fn)
-        async def wrapped(*args, **kwargs):
+        @wrapt.decorator
+        async def wrapped(fn, instance, args, kwargs):
+            token = None
             try:
-                token = None
                 current_trace = current_trace_var.get(None)
                 if current_trace is not None:
                     token = current_trace_var.set(None)
@@ -121,13 +120,13 @@ async def wrapped(*args, **kwargs):
                 if token is not None:
                     current_trace_var.reset(token)
 
-        return wrapped
+        return wrapped(fn)
 
     else:
-        @functools.wraps(fn)
-        def wrapped(*args, **kwargs):
+        @wrapt.decorator
+        def wrapped(fn, instance, args, kwargs):
+            token = None
             try:
-                token = None
                 current_trace = current_trace_var.get(None)
                 if current_trace is not None:
                     token = current_trace_var.set(None)
@@ -137,4 +136,4 @@ def wrapped(*args, **kwargs):
                 if token is not None:
                     current_trace_var.reset(token)
 
-        return wrapped
+        return wrapped(fn)

beeline/test_async.py

@@ -1,3 +1,4 @@
+import inspect
 import unittest
 try:
     # The async functionality uses the contextvars module, added in
@@ -241,6 +242,25 @@ async def decorated_fn():
         self.assertEqual(root_span["span"].id, task0_span["span"].parent_id)
         self.assertEqual(root_span["span"].id, task1_span["span"].parent_id)
 
+    @async_test
+    async def test_async_decorator_preserves_signature(self):
+        @self.beeline.traced("task0")
+        async def task0(a, b, c=1):
+            await asyncio.sleep(0.2)
+
+        @beeline.traced(name="output_integers")
+        def output_integers(n):
+            for i in range(n):
+                yield i
+
+        def task1(a, b, c=1):
+            pass
+
+        self.assertEqual(inspect.getfullargspec(task0).args, ["a", "b", "c"])
+        self.assertEqual(inspect.getfullargspec(output_integers).args, ["n"])
+        self.assertEqual(inspect.getfullargspec(task1).args, ["a", "b", "c"])
+        self.assertEqual(inspect.getfullargspec(task1).args, ["a", "b", "c"])
+
     @async_test
     async def test_traceless_spans_in_other_tasks_should_be_ignored(self):
         """Start a span without first starting a trace in the same task.

beeline/test_suite.py

@@ -21,7 +21,6 @@ def get_test_suite():
             else:
                 ts.addTest(inner_test_group)
         filtered_test_suite.addTest(ts)
-
     return filtered_test_suite
 
 

beeline/trace.py

@@ -15,6 +15,8 @@
 
 from contextlib import contextmanager
 
+import wrapt
+
 from beeline.internal import log, stringify_exception
 
 import beeline.propagation
@@ -410,19 +412,19 @@ def traced_impl(tracer_fn, name, trace_id, parent_id):
     """Implementation of the traced decorator without async support."""
     def wrapped(fn):
         if inspect.isgeneratorfunction(fn):
-            @functools.wraps(fn)
-            def inner(*args, **kwargs):
+            @wrapt.decorator
+            def inner(fn, instance, args, kwargs):
                 inner_generator = fn(*args, **kwargs)
                 with tracer_fn(name=name, trace_id=trace_id, parent_id=parent_id):
                     for value in inner_generator:
                         yield value
-            return inner
+            return inner(fn)
         else:
-            @functools.wraps(fn)
-            def inner(*args, **kwargs):
+            @wrapt.decorator
+            def inner(fn, instance, args, kwargs):
                 with tracer_fn(name=name, trace_id=trace_id, parent_id=parent_id):
                     return fn(*args, **kwargs)
-            return inner
+            return inner(fn)
     return wrapped