Correctly handle modules as request.urlconf -- Thanks to Germán M. Bravo (fixes #166)

2 files changed, 31 insertions, 17 deletions

diff --git a/debug_toolbar/middleware.py b/debug_toolbar/middleware.py
index 7cdf7b0..b1ab737 100644
--- a/debug_toolbar/middleware.py
+++ b/debug_toolbar/middleware.py
@@ -1,13 +1,14 @@
 """
 Debug Toolbar middleware
 """
+import imp
 import thread
 
 from django.conf import settings
+from django.conf.urls.defaults import include, patterns
 from django.http import HttpResponseRedirect
 from django.shortcuts import render_to_response
 from django.utils.encoding import smart_unicode
-from django.conf.urls.defaults import include, patterns
 
 import debug_toolbar.urls
 from debug_toolbar.toolbar.loader import DebugToolbar
@@ -71,26 +72,25 @@ class DebugToolbarMiddleware(object):
         __traceback_hide__ = True
         if self.show_toolbar(request):
 
-            urlconf_name = getattr(request, 'urlconf', settings.ROOT_URLCONF)
-            if urlconf_name not in self._urlconfs:
-
-                import imp
+            urlconf = getattr(request, 'urlconf', settings.ROOT_URLCONF)
+            if isinstance(urlconf, basestring):
+                urlconf = __import__(getattr(request, 'urlconf', settings.ROOT_URLCONF), {}, {}, ['*'])
                 
-                original_urlconf = __import__(getattr(request, 'urlconf', settings.ROOT_URLCONF), {}, {}, ['*'])
+            if urlconf not in self._urlconfs:
                 new_urlconf = imp.new_module('urlconf')
                 new_urlconf.urlpatterns = debug_toolbar.urls.urlpatterns + \
                     patterns('',
-                             ('', include(original_urlconf)),
-                             )
+                        ('', include(urlconf)),
+                    )
                 
-                if hasattr(original_urlconf, 'handler404'):
-                    new_urlconf.handler404 = original_urlconf.handler404
-                if hasattr(original_urlconf, 'handler500'):
-                    new_urlconf.handler500 = original_urlconf.handler500
-
-                self._urlconfs[urlconf_name] = new_urlconf
-
-            request.urlconf = self._urlconfs[urlconf_name] 
+                if hasattr(urlconf, 'handler404'):
+                    new_urlconf.handler404 = urlconf.handler404
+                if hasattr(urlconf, 'handler500'):
+                    new_urlconf.handler500 = urlconf.handler500
+
+                self._urlconfs[urlconf] = new_urlconf
+            
+            request.urlconf = self._urlconfs[urlconf] 
 
             toolbar = DebugToolbar(request)
             for panel in toolbar.panels:
diff --git a/debug_toolbar/tests/tests.py b/debug_toolbar/tests/tests.py
index 355244b..83cc9dc 100644
--- a/debug_toolbar/tests/tests.py
+++ b/debug_toolbar/tests/tests.py
@@ -4,7 +4,6 @@ from debug_toolbar.toolbar.loader import DebugToolbar
 from debug_toolbar.utils.tracking import pre_dispatch, post_dispatch, callbacks
 
 from django.conf import settings
-from django.conf.urls.defaults import patterns
 from django.contrib.auth.models import User
 from django.test import TestCase
 
@@ -118,6 +117,21 @@ class DebugToolbarTestCase(BaseTestCase):
             self.assertEquals(request.urlconf.urlpatterns[0]._callback_str, 'debug_toolbar.views.debug_media')
             self.assertEquals(request.urlconf.urlpatterns[-1].urlconf_name.__name__, 'debug_toolbar.urls')
 
+    def test_request_urlconf_module(self):
+        request = self.request
+        
+        request.urlconf = __import__('debug_toolbar.tests.urls').tests.urls
+        request.META = {'REMOTE_ADDR': '127.0.0.1'}
+        middleware = DebugToolbarMiddleware()
+        
+        with Settings(DEBUG=True):
+            middleware.process_request(request)
+            
+            self.assertFalse(isinstance(request.urlconf, basestring))
+            
+            self.assertTrue(hasattr(request.urlconf.urlpatterns[0], '_callback_str'))
+            self.assertEquals(request.urlconf.urlpatterns[0]._callback_str, 'debug_toolbar.views.debug_media')
+            self.assertEquals(request.urlconf.urlpatterns[-1].urlconf_name.__name__, 'debug_toolbar.tests.urls')
 
 class SQLPanelTestCase(BaseTestCase):
     def test_recording(self):