diff options
| -rw-r--r-- | debug_toolbar/panels/sql.py | 9 | ||||
| -rw-r--r-- | debug_toolbar/utils/tracking/__init__.py | 11 |
2 files changed, 11 insertions, 9 deletions
diff --git a/debug_toolbar/panels/sql.py b/debug_toolbar/panels/sql.py index bc70b97..709a5f5 100644 --- a/debug_toolbar/panels/sql.py +++ b/debug_toolbar/panels/sql.py @@ -13,13 +13,12 @@ from debug_toolbar.panels import DebugPanel from debug_toolbar.utils import render_stacktrace from debug_toolbar.utils.sql import reformat_sql from debug_toolbar.utils.tracking.db import CursorWrapper -from debug_toolbar.utils.tracking import replace_call +from debug_toolbar.utils.tracking import replace_method -# Inject our tracking cursor -@replace_call(BaseDatabaseWrapper.cursor) -def cursor(func, self): - result = func(self) +@replace_method(BaseDatabaseWrapper, 'cursor') +def cursor(original, self): + result = original(self) djdt = DebugToolbarMiddleware.get_current() if not djdt: diff --git a/debug_toolbar/utils/tracking/__init__.py b/debug_toolbar/utils/tracking/__init__.py index 4fbefa5..166b1b5 100644 --- a/debug_toolbar/utils/tracking/__init__.py +++ b/debug_toolbar/utils/tracking/__init__.py @@ -20,18 +20,21 @@ def pre_dispatch(func): return wrapped -def replace_call(func): +def replace_method(klass, method_name): + original = getattr(klass, method_name) + def inner(callback): def wrapped(*args, **kwargs): - return callback(func, *args, **kwargs) + return callback(original, *args, **kwargs) - actual = getattr(func, '__wrapped__', func) + actual = getattr(original, '__wrapped__', original) wrapped.__wrapped__ = actual wrapped.__doc__ = getattr(actual, '__doc__', None) wrapped.__name__ = actual.__name__ - _replace_function(func, wrapped) + setattr(klass, method_name, wrapped) return wrapped + return inner |