ⓘ

Many source code tools, such as linters and formatters, work with CST, a tree-structured representation of source code (like AST, but it also retains nodes such as whitespace and comments). This library is a wrapper around these trees, designed for convenient iterative traversal and node replacement.

It is built on top of libcst.

Table of Contents

• Installation
• Changing nodes
• Filters
• Separating registration from execution
• Context

Installation

You can install cstvis with pip:

pip install cstvis

You can also use instld to quickly try this package and others without installing them.

Changing nodes

The basic workflow is very simple:

• Create a Changer instance.
• Register converter functions with the @<changer object>.converter decorator. Each function takes a CST node as its first argument and returns a replacement node.
• If needed, register filters to prevent changes to certain nodes.
• Iterate over individual changes and apply them as needed.

Let me show you a simple example:

from libcst import Subtract, Add
from cstvis import Changer
from pathlib import Path

# Content of the file:
# a = 4 + 5
# b = 15 - a
# c = b + a # kek
changer = Changer(Path('tests/some_code/simple_sum.py').read_text())

@changer.converter
def change_add(node: Add):
    return Subtract(
        whitespace_before=node.whitespace_before,
        whitespace_after=node.whitespace_after,
    )

for x in changer.iterate_coordinates():
    print(x)
    print(changer.apply_coordinate(x))

#> Coordinate(file=None, class_name='Add', start_line=1, start_column=6, end_line=1, end_column=7, converter_id='__main__:change_add:11')
#> a = 4 - 5
#> b = 15 - a
#> c = b + a # kek
#>
#> Coordinate(file=None, class_name='Add', start_line=3, start_column=6, end_line=3, end_column=7, converter_id='__main__:change_add:11')
#> a = 4 + 5
#> b = 15 - a
#> c = b - a # kek

As you can see in the example, the converter function takes an argument with a type hint. You don’t need to write type-checking if statements because the system determines which node types to convert based on this hint. You can omit the annotation entirely, specify Any, or specify libcst.CSTNode, in which case the converter will be applied to all nodes. If you specify a more specific type, such as libcst.Add, the converter will be applied only to those nodes. You can also specify multiple nodes using the | syntax or Union. Finally, several shortcuts are supported: str -> libcst.SimpleString, int -> libcst.Integer, and float -> libcst.Float.

The key part of this example is the last two lines, where we iterate over the coordinates. What does that mean? This library performs each code change in two stages: identify the coordinates of the change and then apply it. This separation makes it possible to distribute the work across multiple threads or even multiple machines. However, this design also has limitations. If you apply one coordinate change, the resulting code will differ from the original and the remaining coordinates will no longer be valid. You can only apply one change at a time.

Filters

A filter is a special function registered with the @<changer object>.filter decorator. It decides whether a specific CST node should be changed, and returns True if yes, or False if no. As with converters, the filter's type hint determines which nodes it is applied to.

Here is another example (part of the code is omitted):

count_adds = 0

@changer.filter
def only_first(node: Add) -> bool:
    global count_adds
    
    count_adds += 1
    
    return True if count_adds <= 1 else False

for x in changer.iterate_coordinates():
    print(x)
    print(changer.apply_coordinate(x))

#> Coordinate(file=None, class_name='Add', start_line=1, start_column=6, end_line=1, end_column=7, converter_id='__main__:change_add:11')
#> a = 4 - 5
#> b = 15 - a
#> c = b + a # kek

As you can see, now the iteration yields only the first possible change, the rest are filtered out automatically because the filter returns False for them.

Separating registration from execution

In some cases, you may want to separate converter and filter registration from execution. In this case, a special type of objects — Collector — can help you. A collector object has the same decorators as Changer objects, and they can be used in exactly the same way. When creating a Changer object, you can pass a collector object to it:

from cstvis import Collector

collector = Collector()

@collector.converter
def change_add(node: Add):
    return Subtract(
        whitespace_before=node.whitespace_before,
        whitespace_after=node.whitespace_after,
    )

changer = Changer(Path('tests/some_code/simple_sum.py').read_text(), collector=collector)

Context

By default, each converter or filter takes a single argument: the node to which it is applied. However, you can also specify a second argument: the context. The system analyzes the signatures of your functions, detects that they expect a second argument, and passes it to them:

from cstvis import Context

@changer.converter
def change_add(node: Add, context: Context):  # <- The function takes a second argument.
    return Subtract(
        whitespace_before=node.whitespace_before,
        whitespace_after=node.whitespace_after,
    )

The context object has two necessary fields and one useful method:

• coordinate with fields start_line: int, start_column: int, end_line: int, end_column: int and some others. This identifies the current location in the code.
• comment - the comment on the first line of the node, if there is one, without the leading #, or None if there is no comment.
• get_metacodes(key: Union[str, List[str]]) -> List[ParsedComment] - a method that returns a list of parsed comments in metacode format associated with this line of code.

You can also pass an arbitrary dictionary to any decorator in this library; a copy of that dictionary will be passed as a meta attribute of the context object:

from libcst import SimpleString
from cstvis import Changer, Context
from pathlib import Path

# Content of the file:
# a = "old string"

changer = Changer(Path('tests/some_code/simple_string.py').read_text())

@changer.converter(meta={'new_value': '"new string"'})
def change_string(node: SimpleString, context: Context):
    return SimpleString(value=context.meta['new_value'])

for x in changer.iterate_coordinates():
    print(x)
    print(changer.apply_coordinate(x))

#> Coordinate(file=None, class_name='SimpleString', start_line=1, start_column=4, end_line=1, end_column=9, converter_id='__main__:change_add:13')
#> a = "new string"