自定义profile

python中检测代码的执行效率，目前比较常见的办法是profile，其C版本是cProfile，使用起来也很简单。
可惜的是，公司的引擎中并没有包含profile，导致这两个办法都行不通。
于是在周末花了点时间详细阅读profile.py，发现原理很简单。一切基于sys.setprofile这个API。
然后，我写了个简化版本的profile，没有其他库的依赖，并且可以自定义输出格式。
so，上代码，备忘。

#coding=utf-8
"""
@author : leecrest
@time   : 14-3-1 下午8:33
@brief  : 性能分析模拟
"""

import os
import sys
import time
import marshal

g_RootBlock = None
g_CurBlock = None
FilterList = ("_cal_",)
MAX_BLOCK = 50


def GetCurTime():
	return time.clock()


#记录每次调用的信息
class CBlock:
	def __init__(self, parent, func_code):
		"""
		@parent：CBlock，上层调用者
		@func_code：func的基础数据
		"""
		self.m_Code         = func_code     #函数信息
		self.m_CurCall      = 0             #当前调用开始时间
		self.m_CallTime     = 0             #总调用时间
		self.m_CallCount    = 0             #总调用次数
		self.m_Parent       = parent        #父类Blk
		self.m_ChildData    = {}            #子类Blk
		self.m_ChildList    = []


	def __str__(self):
		return "<CBlock Name:%s,CallTime:%f,CallCount:%d,Child:%s>" % (self.Name(),
			self.m_CallTime, self.m_CallCount, self.m_ChildList)


	def Name(self):
		return self.m_Code.co_name


	def Call(self):
		self.m_CurCall = GetCurTime()


	def Return(self):
		global g_CurBlock
		self.m_CallTime += GetCurTime() - self.m_CurCall
		self.m_CurCall = 0
		self.m_CallCount += 1
		g_CurBlock = self.m_Parent


	def AddSub(self, func_code):
		global g_CurBlock
		sName = func_code.co_name
		if not sName in self.m_ChildData:
			self.m_ChildData[sName] = CBlock(self, func_code)
		if not self.m_ChildList or self.m_ChildList[-1] != sName:
			self.m_ChildList.append(sName)

		g_CurBlock = self.m_ChildData[sName]

	def GetData(self):
		fCurTime = self.m_CallTime / self.m_CallCount
		fPercent = 100
		if self.m_Parent and self.m_Parent.m_CallTime > 0:
			fPercent = fCurTime * self.m_Parent.m_CallCount * 100 / self.m_Parent.m_CallTime
		data = {
			"name"      : self.m_Code.co_name,
		    "file"      : self.m_Code.co_filename,
		    "line"      : self.m_Code.co_firstlineno,
		    "time"      : self.m_CallTime,
		    "count"     : self.m_CallCount,
		    "avg"       : fCurTime,
		    "percent"   : fPercent,
		    "child"     : []
		}
		for sName in self.m_ChildList:
			blk = self.m_ChildData[sName]
			data["child"].append(blk.GetData())
		return data


	def Save(self, sFile=None):
		if not sFile:
			sFile = self.m_Code.co_name + ".prof"
		dData = self.GetData()
		dList = []
		if os.path.exists(sFile):
			hFile = open(sFile, "rb")
			dList = marshal.load(hFile)
			hFile.close()
		dList.append(dData)
		iSize = len(dList)
		if iSize > MAX_BLOCK:
			dList = dList[iSize-MAX_BLOCK : ]
		hFile = open(sFile, "wb")
		marshal.dump(dList, hFile)
		hFile.close()



#回调事件分发
def Dispatch(frame, event, ret):
	"""
	@frame：函数信息
	@event：事件名称
	@ret:函数调用的返回值
	"""
	global g_CurBlock, g_RootBlock
	sName = frame.f_code.co_name
	if sName in FilterList:
		return
	#print event, sName, frame.f_back.f_code.co_name,
	#if g_CurBlock:
	#	print g_CurBlock
	if event in ("call", "c_call"):
		if g_CurBlock:
			if event == "c_call" and sName == g_CurBlock.Name():
				return
			g_CurBlock.AddSub(frame.f_code)
		else:
			if not g_RootBlock:
				g_RootBlock = CBlock(None, frame.f_code)
			g_CurBlock = g_RootBlock
		g_CurBlock.Call()
	elif event in ("return", "c_return"):
		if not g_CurBlock:
			return
		if event == "c_return" and sName == g_CurBlock.Name():
			return
		g_CurBlock.Return()
	#elif event == "exception":
	#	pass
	#elif event == "c_exception":
	#	pass


def SaveProfile():
	global g_RootBlock, g_CurBlock
	g_RootBlock.Save()
	g_RootBlock = None
	g_CurBlock = None


#修饰器
def PyProfile(func):
	def _call_(*args, **kwargs):
		sys.setprofile(Dispatch)
		func(*args, **kwargs)
		sys.setprofile(None)
		SaveProfile()
	return _call_

使用极为简单，导入文件后，在需要监控的函数前@PyProfile即可。