title: Modifying Code of Function at Runtime
body: <p>I am trying to write a decorator that adds verbose logging to a function via a decorator (a method would be nice too, but I haven't tried that yet). The motivation behind this is that patching in a one-line add_logs decorator call into a box in production is much easier (and safer) than adding 100 debug lines.</p>
  
  <p>Ex:</p>
  
  <pre><code>def hey(name):
      print("Hi " + name)
      t = 1 + 1
  
      if t &gt; 6:
          t = t + 1
          print("I was bigger")
      else:
          print("I was not.")
      print("t = ", t)
      return t
  </code></pre>
  
  <p>I'd like to make a decorator that will transform this into code that does this:</p>
  
  <pre><code>def hey(name):
      print("line 1")
      print("Hi " + name)
  
      print("line 2") 
      t = 1 + 1
  
      print("line 3")
      if t &gt; 6:
          print("line 4")
          t = t + 1
          print("line 5")
          print("I was bigger")
      else:
          print("line 6")
          print("I was not.")
     print("line 7") 
     print("t = ", t)
     print("line 8")
     return t
  </code></pre>
  
  <p>What I've got so far:</p>
  
  <pre><code>import inspect, ast
  import itertools
  import imp
  
  def log_maker():
       line_num = 1
      while True:
          yield ast.parse('print("line {line_num}")'.format(line_num=line_num)).body[0]
          line_num = line_num + 1
  
  def add_logs(function):
      def dummy_function(*args, **kwargs):
          pass
      lines = inspect.getsourcelines(function)
      code = "".join(lines[0][1:])
      ast_tree = ast.parse(code)
      body = ast_tree.body[0].body
  
      #I realize this doesn't do exactly what I want.
      #(It doesn't add debug lines inside the if statement)
      #Once I get it almost working, I will rewrite this
      #to use something like node visitors
      body = list(itertools.chain(*zip(log_maker(), body)))
      ast_tree.body[0].body = body
      fix_line_nums(ast_tree)
      code = compile(ast_tree,"&lt;string&gt;", mode='exec')
  
      dummy_function.__code__ = code
      return dummy_function
  
  def fix_line_nums(node):
      if hasattr(node, "body"):
          for index, child in enumerate(node.body):
              if hasattr(child, "lineno"):
                  if index == 0:
                      if hasattr(node, "lineno"):
                          child.lineno = node.lineno + 1
                      else:
                          # Hopefully this only happens if the parent is a module...
                          child.lineno = 1
                  else:
                      child.lineno = node.body[index - 1].lineno + 1
              fix_line_nums(child)
  
  @add_logs
  def hey(name):
      print("Hi " + name)
      t = 1 + 1
  
      if t &gt; 6:
          t = t + 1
          print("I was bigger")
      else:
          print("I was not.")
      print("t = ", t)
      return t
  
  if __name__ == "__main__":
      print(hey("mark"))
      print(hey)
  </code></pre>
  
  <p>This produces this error: </p>
  
  <pre><code>Traceback (most recent call last):
    File "so.py", line 76, in &lt;module&gt;
      print(hey("mark"))
  TypeError: &lt;module&gt;() takes no arguments (1 given)
  </code></pre>
  
  <p>which makes sense because compile creates a module and of course modules are not callables. I've tried a hundred different ways of making this work at this point but can't come up with a working solution. Any recommendations? Am I going about this the wrong way?</p>
  
  <p>(I haven't been able to find a tutorial for the ast module that actually modifies code at runtime like this. A pointer to such a tutorial would be helpful as well)</p>
  
  <p>Note: I am presently testing this on CPython 3.2, but a 2.6/3.3_and_up solution would be appreciated. Currently the behavior is the same on 2.7 and 3.3.</p>