👨🏼‍🔬 👨🏾‍🏭 👾 Implementamos conversões de código Python 🕵🏿 🤕 💃🏻

Olá, Habr.

Hoje, oferecemos uma tradução de um artigo que aborda um tópico que não é o mais discutido: compilação de código em Python, a saber: trabalhar com árvore de sintaxe abstrata (AST) e código de bytes. Embora o Python seja uma linguagem interpretada, esses recursos são extremamente importantes do ponto de vista da otimização. Hoje falaremos sobre eles.

Você já se perguntou como exatamente o compilador otimiza seu código para que ele funcione mais rápido? Deseja saber o que é uma árvore de sintaxe abstrata (AST) e para que ela pode ser usada?

Este artigo de revisão descreve como o código Python é convertido em formato de árvore (AST). Depois de criar o AST do seu programa, você pode prosseguir na busca por opções de otimização e transformação para o seu código. No entanto, lembre-se de que otimizar programas Python de maneiras não triviais é extremamente difícil .

Código de programa como uma árvore

Como um computador pode garantir que ele avalie expressões do seu código na ordem correta?
Para fazer isso, ele primeiro refaz o código do programa em uma estrutura de árvore chamada AST.

Ao trabalhar com uma linguagem de programação interpretada (como Python), geralmente é aceito que o intérprete passe pelo seu código e faça tudo o que encontrar, no local, sem converter o código Python em código de máquina de forma alguma. No entanto, na prática, esse esquema de execução causa muitos problemas, o que o torna muito inconveniente.
Tomemos, por exemplo, um problema simples como a prioridade dos operadores. Em uma expressão de exibição 3 + 4 * x , a peça é calculada primeiro4 * x, e somente então 3 podem ser adicionados ao resultado da multiplicação. Talvez você tenha aprendido a precedência dos operadores nas aulas de matemática desenhando essas árvores sob a expressão:

Python usa as regras padrão da notação matemática (multiplicação primeiro, depois adição). Para não confundir nada com a prioridade dos operadores, no Python, primeiro é construída uma árvore como na figura anterior. A operação geral é a adição (na raiz da árvore) e, enquanto o lado esquerdo dessa soma é um número regular, à direita temos o produto. A estrutura de dados resultante é assim:

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

BinOpsignifica Operação Binária (Operação Binária) e indica que em operações como adição e multiplicação - dois operandos. Naturalmente, você não receberá nenhuma adição se a parte certa da expressão não tiver o valor correto. Portanto, você deve primeiro se multiplicar.

Na teoria dos compiladores e linguagens de programação, essa árvore é chamada Abstract Syntax Tree , ou AST, para abreviar. O AST no exemplo acima inclui dois nós BinOp, dois nós Nume um nó Name.

Há um recurso interessante no Python - a capacidade de exibir e exibir diretamente o AST para qualquer programa específico do Python. Tudo o que é necessário é importar um módulo padrãoast, analisando o programa e exibindo o resultado na tela (a propósito, analisar é o processo de conversão do código-fonte do programa na árvore AST).

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

No entanto, você notará que haverá nós e campos adicionais no AST gerados pelo Python, e ele será exibido em uma linha, o que torna à primeira vista mais complicado do que realmente é.

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()))))])

Vamos dividi-lo em nós separados, como da última vez - e reabrir o AST, já no topo, como parte de toda a árvore:

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())
            )
        )
    )
])

Obviamente, o Python "pensa" que a linha que fornecemos para análise é um módulo inteiro. O corpo do módulo é uma lista de todas as instruções contidas nele. A única instrução em nosso exemplo é uma expressão Exprcujo significado é exatamente o que discutimos acima.

Nota: o nó Namepossui um campo adicional ctx(abreviado como "contexto"), que possui um valor Load(). Então, o Python diz que usamos o valor armazenado na variável xe não (re) definimos ou excluímos o nome x. Agora, sua jogada, tente analisar algo como você del xou você x = 123, e você verá como o campo ctxno nó Namemuda para Del()ou Store(), respectivamente.

A propósito: se você instalar o móduloastunparse, a saída AST para a tela pode ficar muito mais bonita e até converter o AST novamente em código Python ativo.

O processo de compilação: o restante

Depois de coletar os programas AST, é possível, em princípio, concluir todo o programa, percorrendo o AST e executando as operações na ordem em que são indicados. No entanto, essa abordagem tem pelo menos duas desvantagens. Em primeiro lugar, o AST pode ocupar uma quantidade relativamente grande de memória, especialmente se contiver informações redundantes. Em segundo lugar, o percurso AST pode levar mais tempo do que o necessário. Em resumo: isso pode ser feito, mas é ineficiente.
O compilador não processa o AST diretamente, mas prepara o bytecode, que é então executado na máquina virtual Python. Embora discutir os detalhes desse processo esteja além do escopo deste artigo, o princípio básico é que o compilador converte o AST em notação polonesa reversa (RPN). Em vez de colocar um operador+entre os operandos esquerdo e direito, colocamos após os dois operandos. No exemplo 3 + 4*xacima, obtemos a sequência 3 4 x * +(e essa notação é especialmente boa porque você pode ver imediatamente a partir da sequência: primeiro você precisa executar a multiplicação e somente depois a adição). Como cada um dos cinco elementos nessa sequência pode, em princípio, ser representado como um único byte, esse código é chamado de código de byte. O Python usa a máquina virtual empilhada para executar esse código com eficiência.

Em outras palavras, o processo de compilação de um programa escrito em Python ocorre em dois estágios. Primeiro, o programa recebido pela entrada é analisado e o resultado é uma árvore de sintaxe abstrata (AST). O compilador passa pelo AST e gera o bytecode. Depois disso, o intérprete Python executa esse bytecode. Tendo adotado a otimização, ela pode ser aplicada no nível AST ou no nível do bytecode. Ambas as opções têm suas próprias vantagens e desvantagens.

Por fim, lembre-se de que, embora o AST seja comum em qualquer implementação do Python, o processo de conversão do AST em bytecode pode ser diferente e, em algumas implementações do Python, digamos, JavaScript, em vez de bytecode, pode ser gerado no estágio intermediário.

Paradigmas de outras linguagens de programação

Nem todas as linguagens de programação usam notação infix, como no Python. Dois exemplos dignos de nota nesse caso são o PostScript, onde o programa é escrito diretamente em notação polonesa reversa e o Lisp, é claro, onde os programas geralmente são escritos em notação polonesa. Então, a nossa expressão do exemplo acima, em Lisp tomaria a seguinte forma: (+ 3 (* 4 x)).

Conversão de nó no AST

Tendo um programa AST, como converter partes individuais desta árvore? Com os convenientes recursos internos do Python.

Se dermos uma olhada no AST e, por exemplo, descobrirmos que os campos lefte os rightnós BinOpsão números (nós Num), podemos executar os cálculos correspondentes com antecedência e substituí-los por um BinOpnó normal Num.

Obviamente, você precisa agir com muito cuidado para não alterar o comportamento do programa, fazendo essas transformações. Por exemplo, em len([a(), b(), c(), d()]), fica claro que o resultado é 4. Mas não podemos substituir toda a expressão do número 4 porque quatro funções a, b, c, dainda foram devidamente invocado.

Novamente, comece com uma otimização simples. Sempre que um nome aparecer no código-fonte de um programa pi, substitua-o pelo valor 3.14159265. O módulo Python astjá fornece as estruturas de dados necessárias para fazer isso: uma classe de conversor NodeTransformerque passa por todos os ASTs e verifica se cada nó pode ser substituído. Por padrão, o método de transformação simplesmente retorna o nó de origem para cada nó, para que obtenhamos o mesmo AST a partir do qual começamos. Mas podemos facilmente substituir o método para nós Name, digamos, para que ele verifique pise é e, em seguida, retorne o nó em Numvez do nó com o nome original ...

	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))

Para que o conversor / otimizador passe pela nossa árvore, é necessário chamar seu método visit, que retornará uma nova árvore alterada.

Infelizmente, é impossível compilar e executar o AST resultante; a razão para isso é um detalhe técnico. Isso ainda não está visível, mas (quase) todos os nós no AST também possuem campos linenoe col_offset. Eles indicam a posição exata de um nó específico no código-fonte. Se você não instalá-los corretamente, o compilador jurará e se recusará a trabalhar.

Então, vamos copiar os campos apropriados do nó de origem Namepara o novo nó Num. Você pode então compilar e executar o AST resultante:

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)

Nota: a função de compilação requer não apenas o código-fonte (em que pode ser um programa em si, ou a linha de AST), mas o nome do arquivo (como pedimos "<string>"), bem como um dos três: "exec", "eval"ou "single".

A necessidade de copiar os campos que descrevem a posição do nó no código-fonte surge com bastante frequência. Portanto, o módulo astpossui uma função dedicada copy_locationapenas para esse fim, e podemos escrever:

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

Por fim, você pode estender o exemplo anterior para que ele efetivamente otimize, ou seja, no nó BinOp. De acordo com a regra de transformação, primeiro devemos transformar / otimizar a esquerda e depois o nó direito como parte do BinOp. Se, como resultado, os nós esquerdo e direito forem números, os cálculos poderão ser realizados no local e substituir o original pelo BinOpresultado numérico da operação.

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))

A propósito, o compilador CPython já está otimizando os nós, BinOpcomo mostrado aqui. O código correspondente é escrito em C e é fornecido em Python / ast_opt.c . Observe: o otimizador CPython é mais universal e funciona não apenas com números, como no exemplo que estamos considerando, mas com diferentes tipos de valores constantes.

Verificando nós no AST

Como garantir que as transformações que fizemos foram corretas? Primeiro, você precisa ignorar completamente o AST e inspecionar todo o programa.

O otimizador apresentado acima continua sendo uma falha séria. O que acontece se você redefinir em algum lugar do programa pi? Imagine algo tão simples e inteligível quanto pi = 4. Nosso otimizador simplesmente substituirá pi no lado esquerdo da expressão pelo valor numérico 3.14159265, e o Python se recusará a compilar porque não pode atribuir nada a um valor literal.

Talvez esse seja exatamente o comportamento que você buscou, tornando pi uma constante verdadeira, que é substituída na compilação e nunca pode ser reatribuída, ou seja, não pode obter um valor diferente. No entanto, isso definitivamente viola a semântica do Python.

Então, o que fazer se quisermos manter a semântica do Python, mas substituir pi sempre que possível? Nesse caso, você primeiro precisa percorrer todo o programa e verificar se o valor de está atribuído em algum lugar pi. Até complicarmos: não recorreremos à substituição de pi se pelo menos um ponto do programa tiver uma atribuição de valor para pi.

Agora usamos o nó visitante, semelhante ao nó conversor descrito acima. Diferentemente do conversor, o visitante não pretende alterar nenhum nó, ele simplesmente passa pelo AST e examina os nós (os visita). Assim, os métodos de visita não retornam nada.

No nosso caso, vamos verificar se o nó refere-se Namea pie faz outra coisa senão colocar o valorpi(lembre-se do campo de contexto 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)

O método generic_visit(node)é chamado pelo visitante para cada nó para o qual não fornecemos um método de visita especializado. Em outras palavras: não existe tal método visit_FunctionDefna classe NodeVisitorque poderíamos chamar usando super(). Em relação às definições de função, devemos chamar um visitante genérico para garantir que todo o corpo da função também seja processado corretamente. Caso contrário, poderíamos ocultar a instrução na função global pie alterar globalmente o valor pi, para que nosso otimizador não notasse nada.

Valores locais em Python

Nosso método, que nos permite determinar se o programador mudou pi, acabou sendo bastante rude. No entanto, o compilador Python age de maneira muito semelhante quando determina quais nomes no escopo de uma função correspondem às variáveis locais. Se uma variável for alterada em algum lugar no escopo da função (e não for explicitamente globalizada, por exemplo, usando a instrução global), essa variável será considerada local em todo o escopo da função.

O exemplo a seguir será executado sem a quarta linha. Mas, embora x = 0a quarta linha nunca seja executada, ainda é considerada uma atribuição ax e, portanto, x se torna uma variável local na escala de toda a função e até na linha 3. É por isso que o Python jura que a variável x na terceira linha ainda não importa.

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

Se você estiver interessado nos detalhes de exatamente como o Python funciona aqui, consulte Python / symtable.c .

Conclusão

No Python, como na maioria das linguagens de programação, um programa específico não é executado diretamente do código-fonte. De fato, a tradução do código-fonte ocorre em dois estágios: primeiro, é feita uma árvore de sintaxe abstrata (AST) e, em seguida, o código de bytes para a máquina virtual empilhada. O Python também fornece vários recursos muito bons para analisar e até transformar o AST de qualquer programa Python específico, após o qual o AST modificado pode ser compilado e executado. Assim, podemos implementar facilmente nossas próprias otimizações.

Claro, simplesmente omiti muitos detalhes aqui. Garantir que sua otimização funcione corretamente em todos os casos e circunstâncias possíveis é uma questão muito trivial. No entanto, o objetivo deste artigo não é informar sobre a otimização pronta para uso na produção, mas fornecer uma idéia básica de como o Python analisa o código do programa, para que você possa aprender como convertê-lo corretamente e otimizá-lo.

Implementamos conversões de código Python