Реализуем преобразования кода на Python

Привет, Хабр.

Сегодня мы предлагаем вам перевод статьи, затрагивающей не самую обсуждаемую тему: компиляцию кода в Python, а именно: работу с абстрактным синтаксическим деревом (AST) и байт-кодом. Притом, что Python является интерпретируемым языком, такие возможности в нем чрезвычайно важны с точки зрения оптимизации. О них мы сегодня и поговорим.

Вы когда-нибудь задумывались, как именно компилятор оптимизирует ваш код, чтобы он работал быстрее? Хотите узнать, что такое абстрактное синтаксическое дерево (AST) и для чего оно может использоваться?

В этой обзорной статье рассказано, как код Python преобразуется в древовидную форму (AST). Соорудив AST вашей программы, вы сможете перейти к поиску возможностей оптимизации и преобразования вашего кода. Однако учтите, что оптимизация программ Python нетривиальными способами исключительно сложна.

Код программы как дерево

Как компьютеру убедиться, что он вычисляет выражения из вашего кода в правильном порядке?
Для этого он сначала переделывает код вашей программы в древовидную структуру под названием AST.

При работе с интерпретируемым языком программирования (таким как Python) принято считать, что интерпретатор проходит по вашему коду и выполняет все, что встретит, прямо на месте, не преобразуя код Python в машинный код каким-либо образом. Однако, на практике такая схема выполнения провоцирует немало проблем, из-за которых получается весьма неудобной.
Возьмем, к примеру, такую простую проблему как приоритет операторов. В выражении вида 3 + 4 * x сначала вычисляется часть 4 * x, и лишь потом можно прибавить 3 к результату умножения. Возможно, на уроках математики вы изучали приоритет операторов, рисуя под выражением вот такие деревья:



В Python применяются стандартные правила математической нотации (сначала умножение, затем сложение). Чтобы ничего не перепутать с приоритетом операторов, в Python сначала строится именно такое дерево, как на предыдущей картинке. Общая операция – это сложение (у корня дерева), и, тогда как левая часть этой суммы представляет собой обычное число, справа мы имеем произведение. Результирующая структура данных выглядит так:

BinOp(
  left  = Num(3),
  op    = Add(),
  right = BinOp(
            left  = Num(4),
            op    = Mult(),
            right = Name('x')
          )
)

BinOp означает Binary Operation (Бинарная операция) и указывает на то, что в таких операциях как сложение и умножение – по два операнда. Естественно, никакого сложения у вас не получится, если в правой части выражения не будет правильного значения – поэтому сначала нужно произвести умножение.

В теории компиляторов и языков программирования такое дерево называется Абстрактное Синтаксическое Дерево, или AST для краткости. В AST в вышеприведенном примере входят два узла BinOp, два узла Num и один узел Name.

Есть в Python приятная фича – возможность непосредственно просматривать и выводить AST для любой конкретной программы Python. Все, что для этого требуется – импортировать стандартный модуль ast, сделать парсинг программы, а затем вывести результат на экран (кстати, парсинг – это и есть процесс преобразования исходного кода программы в дерево AST).

import ast
my_tree = ast.parse("3 + 4*x")
print(ast.dump(my_tree))

Однако вы заметите, что в AST, сгенерированном Python, будут дополнительные узлы и поля, а выводиться оно будет в одну строку, из-за чего на первый взгляд оно кажется сложнее, чем есть на самом деле.

Module(body=[Expr(value=BinOp(left=Num(n=3), op=Add(), right=BinOp(left=Num(n=4), op=Mult(), right=Name(id='x', ctx=Load()))))])

Давайте разобьем его на отдельные узлы, как и в прошлый раз – и заново откроем AST, уже сверху, как часть всего дерева:

Module(body = [
    Expr(
        value = BinOp(
            left  = Num(n=3),
            op    = Add(),
            right = BinOp(
                left  = Num(n=4),
                op    = Mult(),
                right = Name(id='x', ctx=Load())
            )
        )
    )
])

Очевидно, Python «считает», что строка, которую мы дали ему для парсинга – это целый модуль. Тело модуля – это список всех инструкций, содержащихся в нем. Единственная инструкция в нашем примере – это выражение Expr, значение которой именно такое, что мы обсуждали выше.

Обратите внимание: у узла Name есть дополнительное поле ctx (сокращенно «контекст»), имеющее значение Load(). Так Python сообщает, что мы используем значение, сохраненное в переменной x, а не (пере)определяем или удаляем имя x. Теперь ваш ход, попробуйте сами распарсить что-нибудь вроде del x или x = 123, и увидите, как поле ctx в узле Name меняется на Del() или Store(), соответственно.

Кстати: если установить модуль astunparse, то вывод AST на экран можно сделать гораздо красивее, и даже преобразовать AST обратно в действующий код Python.

Процесс компиляции: оставшаяся часть

Собрав AST программы, в принципе возможно выполнить всю программу, проходя по AST и выполняя операции в том порядке, в котором они указаны. Однако, у такого подхода как минимум два недостатка. Во-первых, AST может занимать сравнительно много места в памяти, особенно если содержит избыточную информацию. Во-вторых, обход AST может занимать больше времени, чем необходимо. Короче говоря: сделать так можно, но это неэффективно.
Компилятор не обрабатывает AST напрямую, а готовит байт-код, который затем выполняется на виртуальной машине Python. Хотя, обсуждение деталей этого процесса выходит за рамки данной статьи, базовый принцип заключается в том, что компилятор транслирует AST в обратную польскую нотацию (RPN). Вместо того, чтобы ставить оператор + между левым и правым операндами, мы ставим его после обоих операндов. В примере 3 + 4*x, приведенном выше, получаем последовательность 3 4 x * + (и данная нотация особенно хороша тем, что из последовательности сразу видно: сначала нужно выполнить умножение, а лишь затем – сложение). Поскольку каждый из пяти элементов в этой последовательности в принципе можно представить единственным байтом, такой код называется байт-кодом. Затем Python использует стековую виртуальную машину для эффективного выполнения этого кода.

Иными словами, процесс компиляции программы, написанной на Python, проходит в два этапа. Сначала программа, полученная на вход, проходит парсинг, и в результате получается абстрактное синтаксическое дерево (AST). Затем компилятор проходит по AST и при этом генерирует байт-код. После чего интерпретатор Python выполняет этот байт-код. Взявшись за оптимизацию, ее можно применить или на уровне AST, или на уровне байт-кода. Обоим этим вариантам свойственны собственные достоинства и недостатки.

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

Парадигмы из других языков программирования

Не во всех языках программирования используется инфиксная нотация, как в Python. Два примера, которые стоит отметить в данном случае — PostScript, где программа записывается непосредственно в обратной польской нотации, и Lisp, конечно же, где программы обычно записываются польской нотацией. Так, наше выражение из рассмотренного выше примера, в Lisp приняло бы такой вид: (+ 3 (* 4 x)).

Преобразование узла внутри AST

Имея AST программы, как преобразовывать отдельные части этого дерева? При помощи удобных встроенных возможностей Python.

Если мы взглянем на AST и обнаружим, к примеру, что оба поля left и right узла BinOp являются числами (узлами Num), то сможем произвести соответствующие вычисления заранее, а затем заменить BinOp обычным узлом Num.

Разумеется, нужно действовать очень осторожно, чтобы не изменить поведение программы, занимаясь такими трансформациями. Например, в len([a(), b(), c(), d()]), понятно, что результат будет 4. Но мы не можем заменить все выражение цифрой 4, поскольку четыре функции a, b, c, d все равно должны правильно вызываться.

Опять же, начнем с простой оптимизации. Везде, где в исходном коде программы встречается имя pi, заменим его значением 3.14159265. В модуле Python ast уже предоставляются структуры данных, необходимые, чтобы это сделать: класс-преобразователь NodeTransformer, который проходит через все AST и проверяет для каждого узла, можно ли его заменить. По умолчанию метод трансформации просто возвращает для каждого узла исходный узел, так что у нас получается то же самое AST, с которого мы начинали. Но мы можем с легкостью переопределить метод для узлов Name, скажем, так, чтобы он проверял, а не pi ли это, а затем возвращал узел Num вместо узла с исходным именем…

	import ast
 
class MyOptimizer(ast.NodeTransformer):
 
    def visit_Name(self, node: ast.Name):
        if node.id == 'pi':
            return ast.Num(n=3.14159265)
        return node
 
tree = ast.parse("y = 2 * pi")
optimizer = MyOptimizer()
tree = optimizer.visit(tree)
print(ast.dump(tree))

Чтобы преобразователь/оптимизатор прошел по нашему дереву, необходимо вызвать его метод visit, который затем вернет новое, измененное дерево.

К сожалению, невозможно скомпилировать и запустить результирующее AST, причиной тому одна техническая деталь. Пока это еще не видно, но (почти) все узлы в AST также имеют поля lineno и col_offset. Они указывают точную позицию конкретного узла в исходном коде. Если не установить их как следует, то компилятор станет ругаться и откажется работать.

Итак, давайте скопируем соответствующие поля из исходного узла Name в новый узел Num. Затем можно скомпилировать и выполнить результирующее AST:

import ast
 
class MyOptimizer(ast.NodeTransformer):
 
    def visit_Name(self, node: ast.Name):
        if node.id == 'pi':
            result = ast.Num(n=3.14159265)
            result.lineno = node.lineno
            result.col_offset = node.col_offset
            return result
        return node
 
tree = ast.parse("print(2 * pi)")
optimizer = MyOptimizer()
tree = optimizer.visit(tree)
code = compile(tree, "<string>", "exec")
exec(code)

Обратите внимание: функция compile требует не только исходника (в качестве которого может выступать сама программа, строка или AST), но и имени файла (в качестве которого мы задаем "<string>"), а также одно из трех: "exec", "eval" или "single".

Необходимость скопировать поля, описывающие позицию узла в исходном коде, возникает довольно часто. Поэтому в модуле ast есть выделенная функция copy_location как раз для этой цели, и мы можем написать:

def visit_Name(self, node: ast.Name):
        if node.id == 'pi':
            result = ast.Num(n=3.14159265)
            return ast.copy_location(result, node)
        return node

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

class MyVisitor(ast.NodeTransformer):
 
    def visit_BinOp(self, node: ast.BinOp):
        node.left = self.visit(node.left)
        node.right = self.visit(node.right)
        if isinstance(node.left, ast.Num) and isinstance(node.right, ast.Num):
            if isinstance(node.op, ast.Add):
                result = ast.Num(n = node.left.n + node.right.n)
                return ast.copy_location(result, node)
            elif isinstance(node.op, ast.Mult):
                result = ast.Num(n = node.left.n * node.right.n)
                return ast.copy_location(result, node)
        return node
 
    def visit_Name(self, node: ast.Name):
        if node.id == 'pi':
            result = ast.Num(n=3.14159265)
            return ast.copy_location(result, node)
        return node
 
tree = ast.parse("y = 2 * pi + 1")
optimizer = MyOptimizer()
tree = optimizer.visit(tree)
print(ast.dump(tree))

Кстати, компилятор CPython уже оптимизирует узлы BinOp так, как показано здесь. Соответствующий код написан на C и приводится в Python/ast_opt.c. Обратите внимание: оптимизатор CPython более универсален и работает не только с числами, как в рассматриваемом нами примере, но с различными типами константных значений.

Проверка узлов в AST

Как убедиться, что сделанные нами преобразования были правильными? Для начала нужно полностью обойти AST и осмотреть всю программу.

В оптимизаторе, представленном выше, остается серьезный изъян. Что будет, если где-то в программе вы переопределите pi? Только представьте что-нибудь настолько простое и внятное как pi = 4. Наш оптимизатор просто заменит pi в левой части выражения числовым значением 3.14159265, и Python затем откажется компилироваться, поскольку не может присвоить чего-либо литеральному значению.

Возможно, именно такого поведения вы и добивались, делая pi истинной константной, которая заменяется при компиляции и никогда не может быть переприсвоена, то есть, не может получить другое значение. Однако это, определенно, нарушает семантику Python.

Итак, что же делать, если мы хотим придерживаться семантики Python, но все равно заменить pi везде, где это возможно? В таком случае сначала нужно пройти по всей программе и проверить, не присваивается ли где-нибудь значение для pi. Пока не усложняем: мы не станем прибегать к замене pi, если хоть в одной точке программы есть любое присваивание значения для pi.

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

В нашем случае мы проверяем, ссылается ли узел Name на pi и делает ли что-либо кроме загрузки значения pi (помним о поле контекста ctx).

import ast
 
class MyVisitor(ast.NodeVisitor):
 
    def __init__(self):
        self.modify_pi = False
 
    def visit_FunctionDef(self, node: ast.FunctionDef):
        if node.name == 'pi':
            self.modify_pi = True
        self.generic_visit(node)
 
    def visit_Name(self, node: ast.Name):
        if node.id == 'pi' and not isinstance(node.ctx, ast.Load):
            self.modify_pi = True
 
program = """
def pi():
    return 3.1415
print(2 * pi())
"""
tree = ast.parse(program)
my_visitor = MyVisitor()
my_visitor.visit(tree)
print("Pi modified:", my_visitor.modify_pi)

Метод generic_visit(node) вызывается посетителем для каждого узла, для которого мы не предоставляем специализированный визитирующий метод. Иными словами: нет такого метода visit_FunctionDef в классе NodeVisitor, который мы могли бы вызвать при помощи super(). Что касается определений функций, мы должны вызывать обобщенный посетитель, чтобы убедиться, что и все тело функции также обрабатывается корректно. В противном случае можно было бы спрятать в функции инструкцию global pi и глобально изменить значение pi, так, что наш оптимизатор ничего бы и не заметил.

Локальные значения в Python

Наш метод, позволяющий определить, не изменил ли программист значение pi, получился довольно грубым. Тем не менее, компилятор Python действует весьма схожим образом, когда определяет, какие имена в области видимости функции соответствуют локальным переменным. Если переменная изменяется где-либо в области видимости функции (и явно не сделана глобальной, например, при помощи инструкции global), то эта переменная считается локальной во всей области видимости функции.

Следующий пример отлично выполнится и без четвертой строки. Но, хотя x = 0 в четвертой строке никогда и не выполняется, это все равно считается присваиванием к x и, следовательно, x становится локальной переменной в масштабах всей функции, и даже в строке 3. Вот почему Python будет ругаться, что переменная x в третьей строке еще не имеет значения.

x = 1
def print_x():
    print(x)
    if False: x = 0
print_x()

Если вас интересуют подробности, как именно здесь работает Python, посмотрите Python/symtable.c.

Заключение

В Python, как и в большинстве языков программирования, конкретная программа не исполняется непосредственно из исходного кода. На самом деле, трансляция исходного кода происходит в два этапа: из него сначала делается абстрактное синтаксическое дерево (AST), а затем байт-код для стековой виртуальной машины. В Python также предоставляется ряд очень приятных возможностей для анализа и даже преобразования AST любой конкретной программы Python, после чего измененное AST можно компилировать и выполнять. Таким образом, мы можем с легкостью внедрить наши собственные оптимизации.

Разумеется, немало деталей я здесь просто опустил. Чтобы убедиться, что ваша оптимизация корректно сработает во всех возможных случаях и обстоятельствах – дело весьма нетривиальное. Однако цель этой статьи – не рассказать вам об оптимизации, готовой для использования в продакшене, а дать базовое представление о том, как Python анализирует ваш программный код, чтобы вы научились его правильно преобразовывать, а затем и оптимизировать.