Модифицикация байт-кода функции в Python

Некоторое время назад мне потребовалось решить достаточно необычную задачу, а именно, добавить нестандартный оператор в языке python. Эта задача заключалась в генерации питоновского кода по псевдокоду, похожему на ассемблер, который содержит в себе оператор goto. Сложный лексический анализатор мне писать не хотелось, оператор goto в псевдокоде использовался для организации циклов и переходов по условиям, и хотелось иметь некоторый его аналог в питоне, которого нет.

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

Итак, есть проблема, как добавить в питон пару новых команд и как заставить его их верно интерпретировать (переходить по нужным адресам). Для этого напишем декоратор, который будет подцепляться к функции, в пределах которой мы хотим использовать оператор goto и добавлять метки (label), и воспользуемся модулями dis, который позволяет работать с байт-кодом питона, и new, который позволяет создавать внутренние объекты питона динамически.

Для начала, определимся с форматом команд. Так как питон имеет ряд ограничений по синтаксису, то команды вида

a:
goto a

сделать не удастся. Однако, питон позволяет добавить конструкции вида

label .a
goto .a

Здесь следует заметить, что точка играет важную роль, т.к. питон пропускает пробелы и сводит это к обращениям к атрибутам класса. Запись без точки приведет к сообщению о синтаксической ошибке. Итак, рассмотрим байт-код данных команд. Для этого выполним следующий код:

>>> def f():
>>>     label .a
>>>     goto .a
>>> import dis
>>> dis.dis( f )
  2           0 LOAD_GLOBAL              0 (label)
              3 LOAD_ATTR                1 (a)
              6 POP_TOP

  3           7 LOAD_GLOBAL              2 (goto)
             10 LOAD_ATTR                1 (a)
             13 POP_TOP
             14 LOAD_CONST               0 (None)
             17 RETURN_VALUE

Следовательно, команда объявления метки и перехода по метке сводится к трем операциям LOAD_GLOBAL, LOAD_ATTR, POP_TOP, основные из которых — первые две. Модуль dis позволяет определить байт-код этих команд с помощью словаря opmap и получить по байт-коду их символьное представление с помощью словаря opname.

>>> dis.opmap[ 'LOAD_GLOBAL' ]
116
>>> dis.opmap[ 'LOAD_ATTR' ]
105

Байтовое представление функции f хранится в f.func_code.co_code, а символьные представления ее переменных хранятся в f.func_code.co_names.

>>> f.func_code.co_names
('label', 'a', 'goto')

Теперь немного о байтовых представлениях интересующих нас команд. По куску дизассемблера видно, что команды LOAD_GLOBAL и LOAD_ATTR представляются тремя байтами (слева указано смещение), первый из которых — байт-код операции (из opmap), второй и третий — данные (младший и старший байт соответственно), представляющие собой индекс в списке f.func_code.co_names, соответствующий тому, какую переменную или какой атрибут мы хотим объявить.

Определить, есть ли аргументы у команды (и таким образом, длину команды в байтах), можно с помощью сравнения с dis.HAVE_ARGUMENT. Если она больше или равна данной константе, то она имеет аргументы, иначе — нет. Таким образом, получаем функцию для разбора байт-кода функции. Далее, заменяем код меток на операцию NOP, а код операторов goto на JUMP_ABSOLUTE, которая в качестве параметра принимает смещение внутри функции. Вот, практически и все. Код декоратора и пример использования приведен ниже.

import dis, new

class MissingLabelError( Exception ):
    pass

class ExistingLabelError( Exception ):
    pass
    
def goto( function ):
    labels_dict = {}
    gotos_list = []
    command_name = ''
    previous_operation = ''
    i = 0

    while i < len( function.func_code.co_code ):
        operation_code = ord( function.func_code.co_code[ i ] )
        operation_name = dis.opname[ operation_code ]

        if operation_code >= dis.HAVE_ARGUMENT:
            lo_byte = ord( function.func_code.co_code[ i + 1 ] )
            hi_byte = ord( function.func_code.co_code[ i + 2 ] )
            argument_position = ( hi_byte << 8 ) ^ lo_byte
            
            if operation_name == 'LOAD_GLOBAL':
                command_name = function.func_code.co_names[ argument_position ]
                
            if operation_name == 'LOAD_ATTR' and previous_operation == 'LOAD_GLOBAL':
                if command_name == 'label':
                    label = function.func_code.co_names[ argument_position ]
                    if labels_dict.has_key( label ):
                        raise ExistingLabelError( 'Label redifinition: %s' % label )
                    labels_dict.update( { label : i - 3 } )
                elif command_name == 'goto':
                    gotos_list += [ ( function.func_code.co_names[ argument_position ], i - 3 ) ]
            
            i += 3
            
        else:
            i += 1

        previous_operation = operation_name
        
    codebytes_list = list( function.func_code.co_code )
    for label, index in labels_dict.items():
        codebytes_list[ index : index + 7 ] = [ chr( dis.opmap[ 'NOP' ] ) ] * 7
    # заменяем 7 последовательно идущих байт команд LOAD_GLOBAL, LOAD_ATTR и POP_TOP на NOP

    for label, index in gotos_list:
        if label not in labels_dict:
            raise MissingLabelError( 'Missing label: %s' % label )
        
        target_index = labels_dict[ label ] + 7
        codebytes_list[ index ] = chr( dis.opmap[ 'JUMP_ABSOLUTE' ] )
        codebytes_list[ index + 1 ] = chr( target_index & 0xFF )
        codebytes_list[ index + 2 ] = chr( ( target_index >> 8 ) & 0xFF )

    # создаем байт-код для новой функции
    code = function.func_code
    new_code = new.code( code.co_argcount, code.co_nlocals, code.co_stacksize, code.co_flags,
        str().join( codebytes_list ), code.co_consts, code.co_names, code.co_varnames,
        code.co_filename, code.co_name, code.co_firstlineno, code.co_lnotab )
    
    # создаем новую функцию
    new_function = new.function( new_code, function.func_globals )
    return new_function

Пример использования:

@goto
def test_function( n ):
    
    goto .label1
     
    label .label2
    print n
    goto .label3
        
    label .label1
    print n
    n -= 1
    if n != 0:
        goto .label1
    else:
        goto .label2
        
    label .label3
    print 'the end'

test_function( 10 )

Результат выполнения примера:

10
9
8
7
6
5
4
3
2
1
0
the end

В заключение хочу добавить, что данное решение не вполне соответствует общему стилю питона: оно не слишком надежно из-за сильной зависимости от версии интерпретатора (в данном случае использовался интерпретатор 2.7, но должно работать для всех версий 2-ки), однако решение данной задачи еще раз доказывает большую гибкость языка и возможность добавления новой необходимой функциональности.