2to3 to rewrite testify assertions as plain asserts

rewrite_asserts.py

from lib2to3.fixer_base import BaseFix
from lib2to3.fixer_util import Name
from lib2to3.main import StdoutRefactoringTool
from lib2to3.pytree import Node
import sys


class FixAssertEqual(BaseFix):
    PATTERN = '''
    power<
        'T'
        trailer< '.' 'assert_equal' >
        trailer<
            '(' arglist<
                left=any ',' right=any [ ',' [ message=any ] ]
            > ')'
        >
    >
    '''

    def transform(self, node, results):
        # TODO: this doesn't actually handle the optional third argument
        # TODO: this will produce bad python if the arguments are wrapped
        assert results
        left = results['left'].clone()
        right = results['right'].clone()
        op = Node(self.syms.comp_op, [Name(u'==', prefix=u' ')])
        new = Node(
            self.syms.assert_stmt, [
                Name(u'assert'),
                Node(self.syms.comparison, (left, op, right), prefix=u' '),
            ],
            prefix=node.prefix,
        )
        return new


# TODO so, uh, RefactoringTool only accepts fixers by name, and does name
# mangling to import them.  probably override get_fixers.
FixMain = FixAssertEqual


def main(fn):
    tool = StdoutRefactoringTool(
        ['__main__'],
        options=None,
        explicit=None,
        nobackups=False,
        show_diffs=True,
    )
    tool.refactor_file(fn, write=True)


if __name__ == '__main__':
    main(*sys.argv[1:])

test_file0.py

import testify as T


class TestCase(T.TestCase):
    """this is some kind of test class that tests something"""
    def test_something(self):
        a = 1
        b = 2
        c = a + b
        # we want to make sure addition works here
        T.assert_equal(c, 3)  # see, because 1 + 2 == 3

test_file1.py

import testify as T


class TestCase(T.TestCase):
    """this is some kind of test class that tests something"""
    def test_something(self):
        a = 1
        b = 2
        c = a + b
        # we want to make sure addition works here
        assert c == 3  # see, because 1 + 2 == 3