Merge pull request #16 from yakobu/master

Now we can wrap function with free variables

ipdbugger/__init__.py

@@ -99,10 +99,9 @@ def start_debugging():
     # Get the frame with the error.
     test_frame = sys._getframe(-1).f_back
 
-    from ipdb.__main__ import wrap_sys_excepthook, def_colors
+    from ipdb.__main__ import wrap_sys_excepthook
     wrap_sys_excepthook()
-    IPDBugger(exc_info=sys.exc_info(),
-              color_scheme=def_colors).set_trace(test_frame)
+    IPDBugger(exc_info=sys.exc_info()).set_trace(test_frame)
 
 
 class ErrorsCatchTransformer(ast.NodeTransformer):
@@ -163,8 +162,7 @@ def try_except_handler(self, node):
             self.exception_handlers, new_exception_handlers
 
         # Run recursively on all sub nodes with the new ignore list
-        for item in node.body:
-            self.visit(item)
+        node.body = [self.visit(node_item) for node_item in node.body]
 
         # Revert changes from ignore list
         self.exception_handlers = old_exception_handlers
@@ -301,9 +299,25 @@ def wrapper(*args, **kw):
 
             ast.fix_missing_locations(tree)
 
+            # Define the wrapping function object
+            function_definition = "def _free_vars_wrapper(): pass"
+            wrapping_function = ast.parse(function_definition).body[0]
+
+            # Initialize closure's variables to None
+            body_list = [ast.parse("{var} = None".format(var=free_var)).body[0]
+                         for free_var in victim.__code__.co_freevars]
+
+            # Add the original function ("victim") to the wrapping function
+            body_list.append(tree.body[0])
+
+            wrapping_function.body = body_list
+
+            # Replace original function ("victim") with the wrapping function
+            tree.body[0] = wrapping_function
+
             # Create a new runnable code object to replace the original code
             code = compile(tree, victim.__code__.co_filename, 'exec')
-            victim.__code__ = code.co_consts[0]
+            victim.__code__ = code.co_consts[0].co_consts[1]
 
             # Set a flag to indicate that the method was wrapped
             victim._ipdebug_wrapped = True

setup.py

@@ -1,7 +1,7 @@
 """Setup file for handling packaging and distribution."""
 from setuptools import setup
 
-__version__ = "2.1.0"
+__version__ = "2.2.0"
 
 setup(
     name="ipdbugger",

tests/test_debug.py

@@ -210,3 +210,17 @@ def func():
             pass
 
     func()
+
+
+def test_wrapping_function_with_closure():
+    """Test wrapping function with closure"""
+    raise_exc = True
+
+    @debug
+    def func():
+        if raise_exc:
+            raise ValueError()
+
+    with capture_output(), patch('bdb.Bdb.set_trace') as set_trace:
+        func()
+        assert set_trace.called_once