Added support for class methods.

2 files changed, 18 insertions, 19 deletions

diff --git a/debug_toolbar/utils/tracking/__init__.py b/debug_toolbar/utils/tracking/__init__.py
index 6d4b0e3..766c248 100644
--- a/debug_toolbar/utils/tracking/__init__.py
+++ b/debug_toolbar/utils/tracking/__init__.py
@@ -51,8 +51,7 @@ def _replace_function(func, wrapped):
             module = import_module(func.__module__)
             setattr(module, func.__name__, wrapped)
     elif getattr(func, 'im_self', None):
-        # TODO: classmethods
-        raise NotImplementedError
+        setattr(func.im_self, func.__name__, classmethod(wrapped))
     elif hasattr(func, 'im_class'):
         # for unbound methods
         setattr(func.im_class, func.__name__, wrapped)
diff --git a/tests/tests.py b/tests/tests.py
index c599b6b..5a05050 100644
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -315,23 +315,23 @@ class TrackingTestCase(BaseTestCase):
         self.assertTrue('foo' in foo['kwargs'])
         self.assertEquals(foo['kwargs']['foo'], 'bar')
 
-        # callbacks['before'] = {}
-        #
-        #         @pre_dispatch(TrackingTestCase.class_method)
-        #         def test(**kwargs):
-        #             foo.update(kwargs)
-        #
-        #         self.assertTrue(hasattr(TrackingTestCase.class_method, '__wrapped__'))
-        #         self.assertEquals(len(callbacks['before']), 1)
-        #
-        #         TrackingTestCase.class_method()
-        #
-        #         self.assertTrue('sender' in foo, foo)
-        #         # best we can do
-        #         self.assertEquals(foo['sender'].__name__, 'class_method')
-        #         self.assertTrue('start' in foo, foo)
-        #         self.assertTrue('stop' not in foo, foo)
-        #         self.assertTrue('args' in foo, foo)
+        callbacks['before'] = {}
+
+        @pre_dispatch(TrackingTestCase.class_method)
+        def test(**kwargs):
+            foo.update(kwargs)
+
+        self.assertTrue(hasattr(TrackingTestCase.class_method, '__wrapped__'))
+        self.assertEquals(len(callbacks['before']), 1)
+
+        TrackingTestCase.class_method()
+
+        self.assertTrue('sender' in foo, foo)
+        # best we can do
+        self.assertEquals(foo['sender'].__name__, 'class_method')
+        self.assertTrue('start' in foo, foo)
+        self.assertTrue('stop' not in foo, foo)
+        self.assertTrue('args' in foo, foo)
 
     def test_post_hook(self):
         foo = {}