Некоторое время назад мне потребовалось решить достаточно необычную задачу, а именно, добавить нестандартный оператор в языке 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-ки), однако решение данной задачи еще раз доказывает большую гибкость языка и возможность добавления новой необходимой функциональности.
Автор: dimitrimus