clang_function_tracer.py

#!/usr/bin/env python
# encoding: utf-8

import clang
import clang.cindex
from clang.cindex import CursorKind as CK
import os
import pprint
import sys

class Function(object):
	def __init__(self, cursor, name, params, callees=None, callers=None):
		# self.cursor.location.line
		# self.cursor.location.file.name
		self.cursor = cursor

		self.name = name
		self.params = params
		self.callees = [] if callees is None else callees
		self.callers = [] if callers is None else callers
		self.signature = "{}__{}".format(
			name,
			"_".join(str(param) for param in self.params)
		)

class FunctionCall(object):
	def __init__(self, cursor, target_name, params, calling_function):
		self.cursor = cursor
		self.target_name = target_name
		self.params = params
		self.calling_function = calling_function
	
	def __repr__(self):
		return "<{}({}) from {} at {}:{}>".format(
			self.target_name,
			", ".join(["X"] * len(self.params)),
			self.calling_function,
			self.cursor.location.file.name,
			self.cursor.location.line
		)

class FunctionTracer(object):
	def __init__(self, tu):
		self._tu = tu

		self.cursor_switch = {
			CK.TRANSLATION_UNIT	: self._handle_translation_unit,
			CK.FUNCTION_DECL	: self._handle_function_decl,
			CK.CALL_EXPR		: self._handle_call_expr
		}

		self._level = -1

		self._functions = {}

		self._curr_function = None
		self._handle_cursor(tu.cursor)
	
	def trace_func(self, func_name):
		if func_name not in self._functions:
			raise Exception("No information for func '{}' recorded".format(func_name))

		res = []
		for func in self._functions[func_name]:
			for x in self._do_trace_func(func):
				res.append(x)

		return res
	
	# -----------------------------------------------

	def _do_trace_func(self, func):
		res = []

		for caller in func.callers:
			res.append((caller, self.trace_func(caller.calling_function)))

		return res

	# -----------------------------------------------

	def _add_function(self, func):
		self._functions.setdefault(func.name, []).append(func)

	def _set_curr_func(self, func):
		self._curr_function = func
	
	def _add_func_caller(self, func_call):
		if func_call.target_name not in self._functions:
			raise Exception("Function '{}' not defined?".format(func_call.target_name))

		called_func_opts = self._functions[func_call.target_name]

		for matched_func in self._functions[func_call.target_name]:
			if len(matched_func.params) <> len(func_call.params):
				continue
			matched_func.callers.append(func_call)

		self._curr_function.callees.append(func_call)
	
	def _add_curr_func_callee(self, callee):
		self._curr_function.callees.setdefault(callee.name, []).append(callee)

	# -----------------------------------------------

	def _handle_call_expr(self, cursor):
		params = []
		for param in cursor.get_arguments():
			params.append(param)

		called_func_name = cursor.displayname

		func_call = FunctionCall(cursor, called_func_name, params, self._curr_function.name)
		self._add_func_caller(func_call)

	def _handle_function_decl(self, cursor):
		params = []
		for arg in cursor.get_arguments():
			params.append(arg.displayname)

		func = Function(cursor, cursor.displayname.split("(")[0], params)
		self._set_curr_func(func)
		self._add_function(func)

		for child in cursor.get_children():
			self._handle_cursor(child)

	def _handle_translation_unit(self, cursor):
		for child in cursor.get_children():
			self._handle_cursor(child)

	def _handle_cursor(self, cursor):
		self._level += 1

		#print("{}visiting {}".format("    "*self._level, cursor.kind))

		res = None
		if cursor.kind in self.cursor_switch:
			res = self.cursor_switch[cursor.kind](cursor)
		else:
			for child in cursor.get_children():
				self._handle_cursor(child)

		self._level -= 1
		return res

def print_callstacks(stack, *curr_stack):
	for func_call, callers in stack:
		formatted = "{}:{} in {}".format(
			func_call.cursor.location.file.name,
			func_call.cursor.location.line,
			func_call.calling_function
		)

		if len(callers) == 0:
			print(formatted)
			if len(curr_stack) > 0:
				print("\n".join(curr_stack))
			print("-------------------")

		else:
			print_callstacks(callers, formatted, *curr_stack)
	
if __name__ == "__main__":
	if len(sys.argv) < 3:
		print("USAGE: {} <filename> <functionname>".format(sys.argv[0]))
		exit(1)
	
	filename = sys.argv[1]
	function_name = sys.argv[2]

	clang.cindex.Config.set_library_file("/Library/Developer/CommandLineTools/usr/lib/libclang.dylib")

	index = clang.cindex.Index.create()
	tu = index.parse(filename, args=["-x", "c++"])
	diagnostics = list(tu.diagnostics)

	if len(diagnostics) > 0:
		print("There were parse errors")
		pprint.pprint(diagnostics)
		exit(1)

	tracer = FunctionTracer(tu)
	call_stacks = tracer.trace_func(function_name)
	print_callstacks(call_stacks)

eh, looks like I have an error where it shows the same path multiple times. Oh well, it's just a simple test script.