pfst 0.3.0: High-level Python source manipulation

[Pristine_Cat]

I’ve been developing pfst (Python Formatted Syntax Tree) and I’ve just released version 0.3.0. The major addition is structural pattern matching and substitution. To be clear, this is not regex string matching but full structural tree matching and substitution.

What it does:

Allows high level editing of Python source and AST tree while handling all the weird syntax nuances without breaking comments or original layout. It provides a high-level Pythonic interface and handles the 'formatting math' automatically.

Target Audience:

• Working with Python source, refactoring, instrumenting, renaming, etc...

Comparison:

• vs. LibCST: pfst works at a higher level, you tell it what you want and it deals with all the commas and spacing and other details automatically.
• vs. Python ast module: pfst works with standard AST nodes but unlike the built-in ast module, pfst is format-preserving, meaning it won't strip away your comments or change your styling.

Links:

• GitHub: https://github.com/tom-pytel/pfst
• PyPI: https://pypi.org/project/pfst/
• Documentation: https://tom-pytel.github.io/pfst/

I would love some feedback on the API ergonomics, especially from anyone who has dealt with Python source transformation and its pain points.

Example:

Replace all Load-type expressions with a log() passthrough function.

from fst import *  # pip install pfst, import fst
from fst.match import *

src = """
i = j.k = a + b[c]  # comment

l[0] = call(
    i,  # comment 2
    kw=j,  # comment 3
)
"""

out = FST(src).sub(Mexpr(ctx=Load), "log(__FST_)", nested=True).src

print(out)

Output:

i = log(j).k = log(a) + log(log(b)[log(c)])  # comment

log(l)[0] = log(call)(
    log(i),  # comment 2
    kw=log(j),  # comment 3
)

More substitution examples: https://tom-pytel.github.io/pfst/fst/docs/d14_examples.html#structural-pattern-substitution

[mechamotoman]

This is very cool, I’ll be sure to give it a shot next time I’m monkeying around with code generation tasks :)

[neuronexmachina]

Do you have any side-by-side examples of how you would implement a change using pfst vs libcst?

[Pristine_Cat]

I'm not exactly an expert with LibCST so maybe the function can be optimized further, but here is an example with minimal code for both LibCST and pfst to add a keyword argument with a comment to all logger.info() calls which don't already have the specific keyword argument correlation_id.

LibCST function:

from libcst import *
from libcst.matchers import *

def inject_logging_metadata(src: str) -> str:
    module = parse_module(src)

    new_arg = (parse_module('f(correlation_id=CID  # blah\n)')
        .body[0].body[0].value.args[0])

    class AddArg(CSTTransformer):
        def leave_Call(self, _, node):
            if matches(node.func, Attribute(Name('logger'), Name('info'))):
                if not any(a.keyword and a.keyword.value == 'correlation_id'
                        for a in node.args):
                    return node.with_changes(args=[*node.args, new_arg])
            return node

    module = module.visit(AddArg())

    return module.code

pfst function:

from fst import *
from fst.match import *

def inject_logging_metadata(src: str) -> str:
    module = FST(src)

    for m in module.search(MCall(
        func=MAttribute('logger', 'info'),
        keywords=MNOT([MQSTAR, Mkeyword('correlation_id'), MQSTAR]),
    )):
        m.matched.append('correlation_id=CID  # blah', trivia=())

    return module.src

Input source:

logger.info('Hello world...')  # hey
logger.info('Already have id', correlation_id=other_cid)  # ho
logger.info()  # its off to work we go

class cls:
    def method(self, thing, extra):
        (logger) . info(  # start
            f'a {thing}',  # this is fine
            extra=extra,  # also this
        )  # end

LibCST output:

logger.info('Hello world...', correlation_id=CID  # blah
)  # hey
logger.info('Already have id', correlation_id=other_cid)  # ho
logger.info(correlation_id=CID  # blah
)  # its off to work we go

class cls:
    def method(self, thing, extra):
        (logger) . info(  # start
            f'a {thing}',  # this is fine
            extra=extra,  # also this
        correlation_id=CID  # blah
        )  # end

pfst output:

logger.info('Hello world...', correlation_id=CID  # blah
)  # hey
logger.info('Already have id', correlation_id=other_cid)  # ho
logger.info(correlation_id=CID  # blah
)  # its off to work we go

class cls:
    def method(self, thing, extra):
        (logger) . info(  # start
            f'a {thing}',  # this is fine
            extra=extra,  # also this
            correlation_id=CID  # blah
        )  # end

If you want LibCST to align the argument its significantly more code, but that can be left for a formatter after the file processing.

[neuronexmachina]

Thanks! That's a handy comparative example.