Source Code Batch Changes

Sometimes, existing code needs to be modified in batches, and these operations are too complex to be achieved through the IDE's replacement function.

AST is a good choice to do that.

For example, perform operations on the following code:

Replace class name Old with New
Replace class variable name OLD with NEW
Replace dict({key: value}) with Base(key=value)
Remove empty class
Remove main

elements.py

class OldElements:
    FOO = ''
    OLD_BTN = {'text': 'Lanbao'}
    OLD_CHECKBOX = {'resource_id': 'Lanbao'}
    OLD_BAR = {'xpath': FOO + 'Lanbao'}


class FooElements(OldElements):
    FOO = ' '
    OLD_BTN = {'text': ''}


class EmptyElements:
    ...


if __name__ == '__main__':
    ...

Visiting AST node through NodeTransformer:

transformer.py

from ast import *

import astor


class _BaseTransformer(NodeTransformer):
    def visit_ClassDef(self, node: ClassDef):
        if len(node.body) == 1:
            return None

        if node.name.startswith('Old'):
            node.name = node.name.replace('Old', 'New', 1)

        for base in node.bases:
            if base.id.startswith('Old'):
                base.id = base.id.replace('Old', 'New', 1)

        return node

    def visit_If(self, node: If):
        if node.test.left.id == '__name__':
            return None


class ElementTransformer(_BaseTransformer):
    def visit_ClassDef(self, node: ClassDef):
        if not (node := super().visit_ClassDef(node)):
            return None

        for sub_node in node.body:
            if not isinstance(sub_node, Assign):
                continue
            for target in sub_node.targets:
                if target.id.startswith('OLD'):
                    target.id = target.id.replace('OLD', 'NEW', 1)

                if isinstance(sub_node.value, Dict):
                    keywords = [
                        keyword(arg=key.value, value=value)
                        for key, value in zip(sub_node.value.keys, sub_node.value.values)
                    ]

                    sub_node.value = Call(
                        func=Name(id='Base', ctx=Load()),
                        args=[],
                        keywords=keywords,
                    )

        return node


if __name__ == '__main__':
    tree = astor.parse_file('elements.py')
    et = ElementTransformer()
    modify_tree = et.visit(tree)
    with open('new_elements.py', 'w') as file:
        file.write(unparse(modify_tree))

Execute transformer.py to get the modified code:

new_elements.py

class NewElements:
    FOO = ''
    NEW_BTN = Base(text='Lanbao')
    NEW_CHECKBOX = Base(resource_id=f'{FOO}Lanbao')
    NEW_BAR = Base(xpath=FOO + 'Lanbao')

class FooElements(NewElements):
    FOO = ' '
    NEW_BTN = Base(text='')

Warning

Notice that ast.parse discards all the comments. And if you need to format the code, you may want to explore tools or libraries like autopep8 or black.

Source Code Batch Changes

Comments