Implement self wrapper

datadog_checks_base/datadog_checks/base/utils/tracing.py

@@ -2,6 +2,7 @@
 # All rights reserved
 # Licensed under Simplified BSD License (see LICENSE)
 import functools
+import inspect
 import os
 
 from ..config import is_affirmative
@@ -55,21 +56,38 @@ def traced_wrapper(self, *args, **kwargs):
 
 
 def tracing_method(f, tracer):
-    @functools.wraps(f)
-    def wrapper(*args, **kwargs):
-        service_name = None
-        if args and len(args) > 0 and hasattr(args[0], "name"):
-            service_name = "{}-integration".format(args[0].name)
-        elif f.__name__ == "__init__":
-            # copy the logic that the AgentCheck init method uses to determine the check name
-            name = kwargs.get('name', '')
-            if len(args) > 0:
-                name = args[0]
-            if name:
-                service_name = "{}-integration".format(name)
-
-        with tracer.trace(f.__name__, resource=f.__name__, service=service_name):
-            return f(*args, **kwargs)
+    signature = inspect.signature(f)
+    if signature.parameters.get('self'):
+        @functools.wraps(f)
+        def wrapper(self, *args, **kwargs):
+            service_name = None
+            if hasattr(self, "name"):
+                service_name = "{}-integration".format(self.name)
+            elif f.__name__ == "__init__":
+                # copy the logic that the AgentCheck init method uses to determine the check name
+                name = kwargs.get('name', '')
+                if len(args) > 0:
+                    name = args[0]
+                if name:
+                    service_name = "{}-integration".format(name)
+
+            with tracer.trace(f.__name__, resource=f.__name__, service=service_name):
+                return f(self, *args, **kwargs)
+
+    else:
+        @functools.wraps(f)
+        def wrapper(*args, **kwargs):
+            service_name = None
+            if f.__name__ == "__init__":
+                # copy the logic that the AgentCheck init method uses to determine the check name
+                name = kwargs.get('name', '')
+                if len(args) > 0:
+                    name = args[0]
+                if name:
+                    service_name = "{}-integration".format(name)
+
+            with tracer.trace(f.__name__, resource=f.__name__, service=service_name):
+                return f(*args, **kwargs)
 
     return wrapper