Вадим Великодный Статьи и заметки

Абстрактные синтаксические деревья и генерация LaTeX

14 Jul 2019

В этой статье мы попробуем разобраться, как работать с абстрактным синтаксическим деревом, представляющим код на языке Python, а заодно написать конвертер из Python в язык разметки математических текстов LaTeX. Разумеется, это сложная задача, и мы будем конвертировать только маленькое подмножество — арифметические и логические выражения. Не то, чтобы это было очень полезно с практической точки зрения, но зато поможет разобраться, как работают библиотеки наподобие Numba.

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

В этой статье будет много кода на языке Python и относительно мало объяснений. Надеюсь, что из кода будет понятна суть. Кроме того предполагается, что читатель немного знаком с LaTeX и тем, как в нём записываются формулы.

Синтаксический анализ

Чтобы преобразовать код с одного языка на другой, необходимо разобраться в структуре кода. Такой анализ структуры называется синтаксическим анализом.

Так, например, человеку, который знает Python, ясно, что выражение x + y * z — это сумма x и произведения y * z. Но для компьютера это просто строка из пяти символов (не считая пробелы). Задача синтаксического анализа найти эти самые сумму и произведение. Причём для этого не обязательно знать, как их вычислять, и что это вообще такое. Такая задача — это уже семантический анализ.

Для компьютерных языков практически всегда известны синтаксические правила, по которым записываются команды. Они обычно достаточно просты, однозначны и их не очень много. Есть даже специальная нотация для записи таких правил — так называемая форма Бэкуса — Наура. И, конечно, есть алгоритмы — порой весьма изощрённые — которые по описанию синтаксиса и исходному тексту восстанавливают структуру.

Но мы не будем сами реализовывать эти алгоритмы. В самом деле, зачем? Ведь интерпретатор Python и так уже выполняет анализ. Просто воспользуемся его встроенным синтаксическим анализатором. Благо, для того, чтобы это сделать достаточно модуля ast из стандартной библиотеки.

Название модуля расшифровывается как abstract syntax tree, то есть абстрактное синтаксическое дерево. Дерево — потому что структуру очень удобно представлять в древовидной форме.

Например, выражению x + y * z будет соответствовать дерево

graph TD + --> x + --> * * --> y * --> z

В этом дереве чётко видна иерархия и сразу понятно, какая операция выполняется первой, а какая второй.

Именно так, а не

graph TD * --> + * --> z + --> x + --> y

Это дерево соответствует выражению (x + y) * z.

Когда есть дерево и видна структура, с ним можно выполнять различные манипуляции и генерировать на его основе код на другом языке, цепочки машинных команд и даже текстовые описания.

Чтобы сгенерировать описание для выражения, надо просто «прочитать» его дерево сверху вниз и слева направо. И тогда x + y * z превратится в «сумма x и произведения y и z», а (x + y) * z — в «произведение суммы x и y и z».

Заметим между делом, что такие «прочтения» формул можно записать в виде + x * y z и * + x y z. А это ни что иное, как префиксная нотация для арифметических выражений — форма записи, прикоторой операция записывается перед операндами. Она, как и постфиксная (в которой операция записывается после операндов) не требует скобок для группировки действий. Этот пример хорошо иллюстрирует удобство древовидной структуры как промежуточного представления.

Так как префиксная форма — это лишь «прочтение» дерева, то можно для удобства записи использовать именно её. Рисовать дерево, особенно большое не всегда удобно. И чтобы было более привычно можно для записи использовать функции, ведь функции как раз и записываются перед операндами.

Тогда x + y * z естественным образом превращается в

Add(x, Mul(y, z))

Add — сложение (addition), Mul — умножение (multiplication).

На эту запись можно взглянуть и иначе: у нас есть два объекта: экземпляры классов Add и Mul. В каждом по две ссылки — на левый и правый аргумент.

Мы уже вплотную подобрались к представлению, которое используется в модуле ast в Python.

Модуль ast

Как я уже говорил, ast — это синтаксический анализатор, принимающий на вход строку с исходным текстом на языке Python и возвращающий дерево в виде объектов с атрибутами и ссылками. То есть, он делает большую часть того, что нам требуется.

Это стандартный модуль, так что его легко подключить:

import ast

Функция ast.parse возвращает абстрактное синтаксическое дерево выражения, построенное как объекты, содержащие ссылки друг на друга.

Рассмотрим, например, выражение x = y + 1. Вершина синтаксического дерева — модуль. Основное его свойство — тело, представляющее собой список команд. В нашем случае это всего одна команда присваивания:

tree = ast.parse('x = y + 1')
print(type(tree.body[0]))

выведет на экран

<class '_ast.Assign'>

То есть, верхний узел поддерева, соответствующего присваиванию, имеет тип Assign (присваивание). У этого узла, в свою очередь, есть атрибуты targets и value, позволяющие получить доступ к левой и правой частям команды присваивания.

Конечно, изучать дерево таким способом не очень удобно. Поэтому можно воспользоваться библиотекой astpretty. Устанавливается она как обычно, через pip. Главная её функция — pprint, выводящая дерево в удобочитаемом виде.

import ast, astpretty
tree = ast.parse('x = y + 1')
astpretty.pprint(tree.body[0])

выведет на экран поддерево для первой команды

Assign(
    lineno=1,
    col_offset=0,
    targets=[Name(lineno=1, col_offset=0, id='x', ctx=Store())],
    value=BinOp(
        lineno=1,
        col_offset=4,
        left=Name(lineno=1, col_offset=4, id='y', ctx=Load()),
        op=Add(),
        right=Num(lineno=1, col_offset=8, n=1),
    ),
)

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

Получаем такое дерево:

graph TD Assign --> |targets| Name Name --> |id| x('x') Assign --> |value| BinOp BinOp --> |left| NameY[Name] NameY --> |id| y('y') BinOp --> |op| Add BinOp --> |right| Num Num --> |n| 1(1)

Обработка арифметических выражений

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

Рекурсивный спуск — не самый эффективный метод, но зато самый простой. А наша цель — разобраться, так что для нас он подходит хорошо.

Мы не будем для экономии места обрабатывать все возможные арифметические и логические операции или как-то упрощать их. Ограничимся стандартным набором: +, -, *, / и **. Причём плюс и минус могут быть как унарными, так и бинарными.

Также будем обрабатывать имена переменных и числа.

Важный момент — нам нужно будет расставлять скобки. В дереве их нет, так что расстановка — это наша ответственность. Можно, конечно, ставить их вокруг каждой операции, но (((xy)+z)+w) читать тяжелее, чем xy+z+w. Так что постараемся расставлять их с умом.

Наша главная функция обработки будет называться translate и для простых выражений она имеет вид:

def translate_expresssion(node):
    if isinstance(node, ast.BinOp):
        return translate_bin_op(node)
    elif isinstance(node, ast.UnaryOp):
        return translate_unary_op(node)
    elif isinstance(node, ast.Num):
        return translate_num(node)
    elif isinstance(node, ast.Name):
        return translate_name(node)
    else:
        raise ValueError('Node {} is not supported'.format(type(node)))

Тут всё достаточно прямолинейно. В зависимости от типа узла вызываем соответствующий обработчик.

Обработаем теперь числа и имена. Тут легко: просто возвращаем их не забыв преобразовать число в строку.

def translate_name(node):
    return node.id

def translate_num(node):
    return str(node.n)

Дальше идут унарные операции. Здесь нужно просто поставить знак операции перед выражением. Но есть одно место, где можно ошибиться: если выражение — это сумма (здесь и дальше будем считать разность разновидностью суммы), то нужны скобки. Иначе выражение -(x + y) превратится -x + y, а это совсем не одно и то же. А вот вокруг произведения скобки не нужны.

Возможно, у читателя возник вопрос: «А что если, минус стоит перед выражением вида x * y + z?». Такой случай обрабатывать отдельно нет нужды. Так как у поддерева для этого выражения верхним узлом будет сумма, то есть скобки будут поставлены. Если же верхний узел — произведение (x * (y + z)), то скобки не нужны, так как о скобках вокруг суммы позаботится произведение.

Для добавления скобок воспользуемся вспомогательной функцией add_brackets, добавляющей скобки переменной длины командами LaTeX \left и \right:

def add_brackets(expression):
    return '\\left( {expression} \\right)'.format(
        expression=expression
    )

Нам не раз придётся проверять, что узел соответствует сложению, поэтому сделаем отдельную функцию для этого:

def is_add(node):
    return isinstance(node, ast.BinOp) and (
        isinstance(node.op, ast.Add) or isinstance(node.op, ast.Sub))

А вот и сам обработчик:

def translate_unary_op(node):
    expression = translate(node.operand)
    if isinstance(node.op, ast.UAdd):
        sign = '+'
    elif isinstance(node.op, ast.USub):
        sign = '-'
    else:
        raise ValueError('Node {} is not supported'.format(type(node)))
    if is_add(node.operand):
        expression = add_brackets(expression)
    return '{sign} {expression}'.format(
        sign=sign,
        expression=expression
    )

Теперь самое интересное — бинарные операции. Скобки нужно добавлять только вокруг операндов произведения и для основания степени. Как и в прошлый раз, только если операнд — сложение.

def translate_bin_op(node):
    left = translate(node.left)
    right = translate(node.right)
    if isinstance(node.op, ast.Mult):
        if is_add(node.left):
            left = add_brackets(left)
        if is_add(node.right):
            right = add_brackets(right)
        return '{left} \\cdot {right}'.format(
            left=left,
            right=right
        )
    elif isinstance(node.op, ast.Div):
        return '\\frac{{{left}}}{{{right}}}'.format(
            left=left,
            right=right
        )
    elif isinstance(node.op, ast.Add):
        return '{left} + {right}'.format(
            left=left,
            right=right
        )
    elif isinstance(node.op, ast.Sub):
        return '{left} - {right}'.format(
            left=left,
            right=right
        )
    elif isinstance(node.op, ast.Pow):
        if is_add(node.left):
            left = add_brackets(left)
        return '{{{left}}} ^ {{{right}}}'.format(
            left=left,
            right=right
        )
    else:
        raise ValueError('Node {} is not supported'.format(type(node)))

Проверим, как работает наш транслятор:

tree = ast.parse('-(x + 2) * y * z + 3 * (2 + v) ** 2 / u')
print(translate(tree.body[0].value))

выведет на экран

- \left( x + 2 \right) \cdot y \cdot z + \frac{3 \cdot {\left( 2 + v \right)} ^ {2}}{u}

Это соответствует выражению

\[- \left( x + 2 \right) \cdot y \cdot z + \frac{3 \cdot {\left( 2 + v \right)} ^ {2}}{u}\]

Вроде неплохо.

Сравнения и логические операции

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

Разберёмся, как устроены сравнения.

import ast, astpretty
tree = ast.parse('x > 1 and not 0 < y <= 2')
astpretty.pprint(tree.body[0].value)

выведет на экран

BoolOp(
    lineno=1,
    col_offset=0,
    op=And(),
    values=[
        Compare(
            lineno=1,
            col_offset=0,
            left=Name(lineno=1, col_offset=0, id='x', ctx=Load()),
            ops=[Gt()],
            comparators=[Num(lineno=1, col_offset=4, n=1)],
        ),
        UnaryOp(
            lineno=1,
            col_offset=10,
            op=Not(),
            operand=Compare(
                lineno=1,
                col_offset=14,
                left=Num(lineno=1, col_offset=14, n=0),
                ops=[
                    Lt(),
                    LtE(),
                ],
                comparators=[
                    Name(lineno=1, col_offset=18, id='y', ctx=Load()),
                    Num(lineno=1, col_offset=23, n=2),
                ],
            ),
        ),
    ],
)

Здесь BoolOp — это цепочка логических операций, а Compare — это цепочка сравнений. Логическое «не» мы добавим к уже имеющимся унарным операциям.

Со скобками поступим просто. У операндов сравнений их нет, а у операндов логических операций есть всегда, если операнд не имя и не число. Исключение — not, которую будем изображать в виде черты над выражением. Выражение под чертой тоже не будем брать в скобки.

Сперва модифицируем translate:

def translate(node):
    if isinstance(node, ast.BinOp):
        return translate_bin_op(node)
    elif isinstance(node, ast.UnaryOp):
        return translate_unary_op(node)
    elif isinstance(node, ast.Num):
        return translate_num(node)
    elif isinstance(node, ast.Name):
        return translate_name(node)
    elif isinstance(node, ast.BoolOp):
        return translate_bool_op(node)
    elif isinstance(node, ast.Compare):
        return translate_compare(node)
    else:
        raise ValueError('Node {} is not supported'.format(type(node)))

Теперь добавим операцию not к унарным:

def translate_unary_op(node):
    expression = translate(node.operand)
    if isinstance(node.op, ast.Not):
        return '\\overline{{{expression}}}'.format(
            expression=expression
        )

    if isinstance(node.op, ast.UAdd):
        sign = '+'
    elif isinstance(node.op, ast.USub):
        sign = '-'
    else:
        raise ValueError('Node {} is not supported'.format(type(node)))

    if is_add(node.operand):
        expression = add_brackets(expression)

    return '{sign} {expression}'.format(
        sign=sign,
        expression=expression
    )

Обработаем сравнения. Нам потребуется вспомогательная функция, которая по типу операции возвращает значок.

COMPARE_OPS = {
    ast.Lt: '<',
    ast.LtE: '\\leq',
    ast.Gt: '>',
    ast.GtE: '\\geq',
    ast.Eq: '=',
}

def compare_op_sign(op):
    try:
        return COMPARE_OPS[type(op)]
    except KeyError:
        raise ValueError('Operator {} is not supported'.format(type(op)))


def translate_compare(node):
    args = [translate(node.left)]
    for op, arg in zip(node.ops, node.comparators):
        args.append(compare_op_sign(op))
        args.append(translate(arg))
    return ' '.join(args)

И, наконец, логические операции.

def bool_op_sign(op):
    if isinstance(op, ast.And):
        return '\\wedge'
    elif isinstance(op, ast.Or):
        return '\\vee'
    else:
        raise ValueError('Operator {} is not supported'.format(type(op)))


def translate_bool_op(node):
    sign = ' {sign} '.format(
        sign=bool_op_sign(node.op)
    )
    args = []
    for arg in node.values:
        expression = translate(arg)
        if not isinstance(arg, ast.Name) and \
            not isinstance(arg, ast.Num) and \
            not (
                isinstance(arg, ast.UnaryOp) and \
                type(arg.op) == ast.Not):
            expression = add_brackets(expression)
        args.append(expression)
    return sign.join(args)

Осталось проверить.

tree = ast.parse('not A >= B and 0 < C <= 1')
print(translate(tree.body[0].value))

Получаем код

\overline{A \geq B} \wedge \left( 0 < C \leq 1 \right)

Это соответствует формуле

\[\overline{A \geq B} \wedge \left( 0 < C \leq 1 \right)\]

Всё, как и планировалось.

Трансляция функций и выполнение

До сих пор мы транслировали не код на Python, строку с кодом на Python. То есть мы не могли передать в наш конвертер уже имеющуюся функцию. Но, к счастью, в Python есть стандартный модуль inspect, позволяющий для любой функции получить её исходный код с помощью функции inspect.getsource. Так можно сделать конвертер, принимающий именно имя функции как аргумент.

В нашем случае, конечно, пользы в этом мало, так как мы не можем в Python заморозить выражение, чтоб оно не было вычислено. Но можно, например, поместить его в лямбда-функцию без аргументов. Правда, при этом придётся обрабатывать и функцию.

Стоит отметить также, что дерево можно скомпилировать с помощью стандартной функции compile, а потом выполнить с помощью exec. Это открывает огромный просто для применения рассмотренного подхода.

Что дальше?

Для языка разметки LaTeX есть замечательный пакет clrscode3e, позволяющий описывать алгоритмы. Он используется для оформления псевдокода из книги Кормена, Лейзерсона, Ривеста и Штайна «Алгоритмы: построение и анализ» (отсюда и название пакета — первые буквы фамилий авторов). Оформление получается красивым, но вот писать команды LaTeX не так удобно, они здесь многословные. Потому мне и пришла в голову мысль попробовать написать простой конвертер из Python (который и сам похож на псевдокод) в LaTeX, которая и подтолкнула к написанию этой статьи.

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

Если же времени нет, но хочется попрактиковаться в работе с деревьями, то можно добавить, например, обработку математических констант и функций из библиотеки math. Кстати, вызовов функций у нас тоже нет.

Если хочется окунуться поглубже, то можно написать код, который трансформирует дерево для некоторой функции. Таким образом, кстати, можно не только генерировать LaTeX, но «добавить» в язык новую команду, оптимизировать код или даже выполнить JIT-компиляцию.

JIT-компиляция — это как раз то, что делает библиотека Numba. Она по исходному коду функции на Python восстанавливает синтаксическое дерево и преобразует его в машинные команды с помощью библиотеки LLVM. Это значительно ускоряет вычисления. С другой стороны, по понятным причинам не любой код на Python можно так преобразовать. Всё же Python очень высокоуровневый, и не всё в нём можно легко оптимизировать.

Вы можете скачать код к этой статье или запустить его в Binder.

Смотрите также