Редактирование кода функции во время выполнения

Question

Я пытаюсь написать декоратор, который добавляет подробное ведение журнала к функции через декоратор (метод тоже был бы хорош, но я еще не пробовал это). Мотивация заключается в том, что исправление в однострочном вызове декоратора add_logs в поле в производстве намного проще (и безопаснее), чем добавление 100 строк отладки.

Ex:

def hey(name): 
    print("Hi " + name) 
    t = 1 + 1 

    if t > 6: 
     t = t + 1 
     print("I was bigger") 
    else: 
     print("I was not.") 
    print("t = ", t) 
    return t 

Я хотел бы сделать декоратора, который преобразует это в код, который делает это:

def hey(name): 
    print("line 1") 
    print("Hi " + name) 

    print("line 2") 
    t = 1 + 1 

    print("line 3") 
    if t > 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 

То, что я получил до сих пор:

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,"<string>", 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 > 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) 

Произошла ошибка:

Traceback (most recent call last): 
    File "so.py", line 76, in <module> 
    print(hey("mark")) 
TypeError: <module>() takes no arguments (1 given) 

, что имеет смысл, поскольку компиляция создает модуль и, конечно, модули не являются вызывающими. Я пробовал сотню различных способов сделать эту работу на данный момент, но не могу найти рабочего решения. Любые рекомендации? Неужели я ошибаюсь?

(я не смог найти учебник для модуля Аст, который фактически изменяет код во время выполнения, как это указатель на такой учебник будет полезным.)

Примечания: В настоящее время я тестирование это на CPython 3.2, но будет оценено решение 2.6/3.3_and_up. В настоящее время поведение одинаково на 2.7 и 3.3.

Answer 1

При компиляции исходного кода, вы получите код объекта, представляющий модуль, а не функция.Подстановка этого кодового объекта в существующую функцию не будет работать, потому что это не объект кода функции, это объект кода модуля. Это все еще объект кода, хотя и не настоящий модуль, поэтому вы можете просто сделать hey.hey, чтобы получить от него функцию.

Вместо этого, как описано в this answer, вам необходимо использовать exec для выполнения кода модуля, сохранения результирующих объектов в словаре и извлечения того, который вы хотите. Что вы можете сделать, грубо говоря, является:

code = compile(ast_tree,"<string>", mode='exec') 
mod = {} 
exec(code, mod) 

mod['hey'] Теперь модифицированная функция. Вы можете переназначить глобальное значение hey или заменить его объект кода.

Я не уверен, что то, что вы делаете с АСТ, совершенно правильно, но вам нужно будет сделать что-то в любом случае, и если в работе АСТ возникнут проблемы, сделайте это, где вы можете начать отлаживать их.

Answer 2

Когда вы вызываете inspect.getsource() с украшенной функцией, вы также получаете декоратор, который в вашем случае вызывается рекурсивно (всего два раза, а во второй раз он производит OSError).

Вы можете использовать этот обходной путь, чтобы удалить @add_logs линию от источника:

lines = inspect.getsourcelines(function) 
code = "".join(lines[0][1:]) 

EDIT:

Похоже, ваша проблема в том, что ваш dummy_function не принимает аргументы:

>>> print(dummy_function.__code__.co_argcount) 
0 
>>> print(dummy_function.__code__.co_varnames) 
() 

Принимая во внимание, что ваша оригинальная функция:

>>> print(hey.__code__.co_argcount) 
1 
>>> print(hey.__code__.co_varnames) 
('name') 

EDIT:

Вы правы объект code возвращается в качестве модуля. Как указано в другом ответе, вы должны выполнить этот объект, а затем назначить полученную функцию (идентифицируемую на function.__name__) на dummy_function.

Как так:

code = compile(ast_tree,"<string>", mode='exec') 
mod = {} 
exec(code, mod) 
dummy_function = mod[function.__name__] 
return dummy_function 

Тогда:

>>> print(hey('you')) 
line 1 
Hi you 
line 2 
line 3 
I was not. 
line 4 
t = 2 
line 5 
2 

Answer 3

Похоже, вы пытаетесь внедрить функцию трассировки. Могу ли я предложить использовать sys.settrace, чтобы сделать это более многоразовым способом?

import sys 

def trace(f): 
    _counter = [0] #in py3, we can use `nonlocal`, but this is compatible with py2 
    def _tracer(frame, event, arg): 
     if event == 'line': 
      _counter[0] += 1 
      print('line {}'.format(_counter[0])) 
     elif event == 'return': #we're done here, reset the counter 
      _counter[0] = 0 
     return _tracer 
    def _inner(*args, **kwargs): 
     try: 
      sys.settrace(_tracer) 
      f(*args, **kwargs) 
     finally: 
      sys.settrace(None) 
    return _inner 

@trace 
def hey(name): 
    print("Hi " + name) 
    t = 1 + 1 

    if t > 6: 
     t = t + 1 
     print("I was bigger") 
    else: 
     print("I was not.") 
    print("t = ", t) 
    return t 

hey('bob') 

Выход:

$ python3 test.py 
line 1 
Hi bob 
line 2 
line 3 
line 4 
I was not. 
line 5 
t = 2 
line 6 

Обратите внимание, что семантика этого немного иначе, чем в вашей реализации; Например, ветви if, не выполняемые вашим кодом, не учитываются.

В результате получается менее хрупким - вы фактически не изменяете код функций, которые вы украшаете, и имеет дополнительную полезность. Функция трассировки дает вам доступ к объекту фрейма перед выполнением каждой строки кода, поэтому вы можете регистрировать locals/globals (или делать некоторые изворотливые инъекции, если вы так склонны).