Automatically add logging to top-level functions with cmonster

add_logging.py

import cmonster
import cmonster.ast
import sys

parser = cmonster.Parser(
"test.cpp",
data="""\
#include <stdio.h>
int main(int argc, const char *argv[])
{
    if (argc % 2 == 0)
    {
        return 1;
    }
    else
    {
        return 0;
    }
}
""")

class Visitor:
    def visit(self, thing):
        method = "visit_%s" % thing.__class__.__name__
        if hasattr(self, method):
            getattr(self, method)(thing)

class FunctionVisitor(Visitor):
    def __init__(self, func_decl, rewriter):
        Visitor.__init__(self)
        self.func_decl = func_decl
        self.rewriter = rewriter

    def visit_IfStatement(self, ifstmt):
        self.visit(ifstmt.then)
        self.visit(ifstmt.else_)

    def visit_CompoundStatement(self, compound):
        for stmt in compound.body:
            self.visit(stmt)

    def visit_ReturnStatement(self, stmt):
        text = 'printf("Returning from %s at line %d\\n");\n' % \
                   (self.func_decl.name, stmt.location.line)
        self.rewriter.insert(stmt, text)


result = parser.parse()
rewriter = cmonster.Rewriter(result)
for decl in result.translation_unit.declarations:
    if decl.location.in_main_file and \
       isinstance(decl, cmonster.ast.FunctionDecl):
        insertion_loc = decl.body[0]
        rewriter.insert(insertion_loc,
            'printf("Entering %s at line %d\\n");\n' % \
                (decl.name, decl.location.line))
        visitor = FunctionVisitor(decl, rewriter)
        for stmt in decl.body:
            visitor.visit(stmt)

rewriter.dump(sys.stdout)