Merge "Add fix for static and class methods in @trace_cls"1.0.0

2 files changed, 48 insertions, 2 deletions

diff --git a/osprofiler/profiler.py b/osprofiler/profiler.py
index 002080c..f7a880e 100644
--- a/osprofiler/profiler.py
+++ b/osprofiler/profiler.py
@@ -148,6 +148,8 @@ def trace_cls(name, info=None, hide_args=False, trace_private=False):
     """
 
     def decorator(cls):
+        clss = cls if type(cls) is type else cls.__class__
+        mro_dicts = [c.__dict__ for c in inspect.getmro(clss)]
         for attr_name, attr in inspect.getmembers(cls):
             if not (inspect.ismethod(attr) or inspect.isfunction(attr)):
                 continue
@@ -156,8 +158,18 @@ def trace_cls(name, info=None, hide_args=False, trace_private=False):
             if not trace_private and attr_name.startswith("_"):
                 continue
 
-            setattr(cls, attr_name,
-                    trace(name, info=info, hide_args=hide_args)(attr))
+            wrapped_obj = None
+            for cls_dict in mro_dicts:
+                if attr_name in cls_dict:
+                    wrapped_obj = cls_dict[attr_name]
+                    break
+
+            wrapped_method = trace(name, info=info, hide_args=hide_args)(attr)
+            if isinstance(wrapped_obj, staticmethod):
+                wrapped_method = staticmethod(wrapped_method)
+            elif isinstance(wrapped_obj, classmethod):
+                wrapped_method = classmethod(wrapped_method)
+            setattr(cls, attr_name, wrapped_method)
         return cls
 
     return decorator
diff --git a/osprofiler/tests/test_profiler.py b/osprofiler/tests/test_profiler.py
index ca38e10..7f86f8d 100644
--- a/osprofiler/tests/test_profiler.py
+++ b/osprofiler/tests/test_profiler.py
@@ -227,6 +227,13 @@ class FakeTracePrivate(FakeTracedCls):
     pass
 
 
+@profiler.trace_cls("rpc")
+class FakeTraceStatic(FakeTracedCls):
+    @staticmethod
+    def method4(arg):
+        return arg
+
+
 def py3_info(info):
     # NOTE(boris-42): py33 I hate you.
     info_py3 = copy.deepcopy(info)
@@ -327,6 +334,33 @@ class TraceClsDecoratorTestCase(test.TestCase):
                       possible_mock_calls("rpc", expected_info))
         mock_stop.assert_called_once_with()
 
+    @mock.patch("osprofiler.profiler.stop")
+    @mock.patch("osprofiler.profiler.start")
+    def test_static(self, mock_start, mock_stop):
+        fake_cls = FakeTraceStatic()
+
+        self.assertEqual(25, fake_cls.method4(25))
+
+        expected_info = {
+            "function": {
+                # fixme(boris-42): Static methods are treated differently in
+                #                  Python 2.x and Python 3.x. So in PY2 we
+                #                  expect to see method4 because method is
+                #                  static and doesn't have reference to class
+                #                  - and FakeTraceStatic.method4 in PY3
+                "name":
+                    "osprofiler.tests.test_profiler.method4" if six.PY2 else
+                    "osprofiler.tests.test_profiler.FakeTraceStatic.method4",
+                "args": str((25,)),
+                "kwargs": str({})
+            }
+        }
+
+        self.assertEqual(1, len(mock_start.call_args_list))
+        self.assertIn(mock_start.call_args_list[0],
+                      possible_mock_calls("rpc", expected_info))
+        mock_stop.assert_called_once_with()
+
 
 @six.add_metaclass(profiler.TracedMeta)
 class FakeTraceWithMetaclassBase(object):