Durante a EuroPython 2012, depois do meu treinamento e das minhas palestras, eu realmente precisava codar, então comecei a hackear em um aplicativo “prático” da AST – tornando (algumas) constantes mais rápidas do que as variáveis locais, em vez de mais lento por tê-las alinhado no import.

Para fazer isso, eu uso um importador personalizado em meta_path para interceptar a importação, e encontrar as atribuições que se assemelham a constantes – isto é, qualquer variável de nível de módulo em ALL_CAPS que seja uma simples string ou um número.

Eu armazeno os nomes e os valores, e depois simplesmente substituo qualquer tentativa de carregar o nome com o valor armazenado. (Sim – o que pode dar errado de várias formas terríveis!)

Para aqueles que não conhecem AST – Abstract Syntax Tree – é uma representação em árvore da estrutura sintática de seu programa. O Python permite que você olhe e modifique a AST antes de compilá-la, o que na prática permite a você reescrever a própria estrutura de um programa antes que ele seja executado. Para uma boa introdução para a AST, eu recomendo What would you do with an AST, de Matthew Desmarais..

Para fazer isso, primeiro eu preciso interceptar a importação – o importador em si não é muito interessante, mas o que ele faz é tentar encontrar a fonte para qualquer módulo importado (se, por exemplo, um arquivo .pyc for encontrado, ele simplesmente retira o ‘c’ e tenta carregar o arquivo.).

Com a fonte encontrada e lida, o importador só chama o transformador, compila o resultado, e coloca-o  no ssys.modules:

module = types.ModuleType(name)
inlined = transform(src)
code = compile(inlined, filename, 'exec')
sys.modules[name] = module
exec(code,  module.__dict__)

O método transform analisa a fonte, cria o NodeTransformer que irá modificar a AST, e transfere a AST analisada para ele.

def transform(src):
    """Transforms the given source and return the AST"""
    tree = ast.parse(src)
    cm = ConstantMaker()
    newtree = cm.visit(tree)
    return newtree

Nosso NodeTransformer é igualmente simples e sobrecarrega o visit_Module (para encontrar as constantes) e o visit_Name (para substituir o uso dos nomes com o valor). O visit_Module começa com a construção de uma lista de todas as atribuições no corpo do módulo, e depois filtra as que cumprem os critérios de constantes: eles devem ser números ou strings, e devem ser nomeados em ALL_CAPS. Quaisquer atribuições são armazenadas em um mapa name->value, que podem então ser utilizadas por visit_Name.

def visit_Module(self, node):
    """Find eglible variables to be inlined and store
    the Name->value mapping in self._constants for later use"""
    assigns = [x for x in node.body if
               type(x) == _ast.Assign]
    for assign in assigns:
        if type(assign.value) in (_ast.Num, _ast.Str):
            for name in assign.targets:
                if RE_CONSTANT.match(name.id):
                    self._constants[name.id] = assign.value
    return self.generic_visit(node)

A análise das atribuições deve ser feita antes da chamada para generic_visit, ou nós não vamos ter o mapeamento até depois de o resto do módulo já ter sido visitado. O mapeamento torna o trabalho de visit_Name extremamente simples:

def visit_Name(self, node):
    """If node.id is in self._constants, replace the
    loading of the node with the actual value"""
    return self._constants.get(node.id, node)

E isso é tudo o que precisamos fazer! Um benchmark simples (simplista?) mostra que ela funciona conforme esperado para casos simples – dada a seguinte fonte que mistura acesso constante com algum “trabalho” de outros:

ONE = 1
TWO = "two"
THREE = 3
FOUR = 4.0
def costly(iterations):
   tempstr = ""
   obj = MyClass()
   for i in range(iterations):
     tempstr += ONE * TWO * THREE + str(i)
     obj.change()
   tempstr += str(obj.value)
    eturn tempstr
class MyClass(object):
   def __init__(self):
     self.value = random.random()*THREE
   def change(self):
     self.value += random.random()*FOUR

…uma versão transformada executa de 15% a 20% mais rápido do que a versão não transformada. Claro, meu primeiro benchmark, que fez apenas o carregamento de constantes, foi zilhões de vezes mais rápido, mas também não muito interessante.

Esta é, naturalmente, uma implementação muito limitada – uma implementação “boa” teria que evitar escrever para as constantes (neste momento, a escrita será silenciosamente ignorada pelo código no módulo atual), a escrita in module para uma constante deve ser detectada, a transformação deveria fazer um fallback para retornar a árvore não-transformada, se falhar e talvez, apenas talvez, não será apenas uma ideia muito boa mesmo.

Foi, no entanto, muito divertido escrever! O código está disponível no repositório do site – o timer que usei para os benchmarks foi escrito por @BaltoRouberol.

A próxima experiência será o alinhamento de funções, eu acho, ou talvez avaliação preguiçosa de parâmetros de função.

***

Texto original disponível em http://blaag.haard.se/Using-the-AST-to-hack-constants-into-Python/#disqus_thread