从代码计时想到的Python计时模块、上下文管理器、变量生存期

Setup

前一篇文章中使用了代码计时的功能，使用time模块获得起始和终止时间然后相减获得运行时间。

import time

start = time.time()
# code to run
end = time.time()
cost = end - start #time in second

Python中除了这种最原始的计时方案，还提供了专门的计时模块timeit，该模块提供了较为完善的计时功能，而且支持重复测试，更多的教程中提倡使用这种方案。以下例子来自博客

def test1():
    n=0
    for i in range(101):
        n+=i
    return n

def test2():
    return sum(range(101))

def test3():
    return sum(x for x in range(101))

if __name__=='__main__':
    from timeit import Timer
    t1=Timer("test1()","from __main__ import test1")
    t2=Timer("test2()","from __main__ import test2")
    t3=Timer("test3()","from __main__ import test3")
    print t1.timeit(1000000)
    print t2.timeit(1000000)
    print t3.timeit(1000000)
    print t1.repeat(3,1000000)
    print t2.repeat(3,1000000)
    print t3.repeat(3,1000000)

以上是两种较为常见的计时功能，有相应的优缺点：

• 第一种方案较为直观，需要了解time模块的使用方法，并且需要进行单位转换
• 第二种方案有官方支持，但是不方便用于测量非函数代码块（例如某几行代码的运行时间），而且需要通过参数来配置测试环境

基于以上考虑，我想能不能实现一种更加灵活的代码计时工具。

上下文管理器 (Context Manager)

初见这个概念可能有点陌生，但下面的代码大家一定再熟悉不过了：

with open('afile.txt', 'r') as f:
    content = f.read()

上下文管理器是Python 2.5之后引入的概念，也就是常用的with语句。上下文管理器大大减轻了操作一些需要打开和关闭的资源（这里找不到好的定语了，就是文件和数据库这种资源）的代码。更重要的一点是，with语句使用很灵活，内部可以包含任意的代码。
实现上下文管理器有两种方法，分别是在自定义类中实现__enter()__ __exit()__和使用装饰器@contextlib.contextmanager

变量生存期

Python对于变量生存期的处理很神奇，这里不讨论太多的内容，只指出两点较为特别的地方：

• 上例中content是在with中定义的变量，但是在外部依旧可以使用
• Python中的循环变量在循环结束之后依旧可以访问到:

for i in range(10):
    pass
print(i) # 9

timeme

基于以上讨论，实现一个用于代码计时的上下文管理器，代码已托管到gist，之后只在gist更新，文中代码不再更新。
这个模块仅仅用于学习上下文管理器和代码计时，不恰当的地方欢迎大家指正。

__author__ = 'mengpeng'
import time
 
 
class timeme(object):
    __unitfactor = {'s': 1,
                    'ms': 1000,
                    'us': 1000000}
 
    def __init__(self, unit='s', precision=4):
        self.start = None
        self.end = None
        self.total = 0
        self.unit = unit
        self.precision = precision
 
    def __enter__(self):
        if self.unit not in timeme.__unitfactor:
            raise KeyError('Unsupported time unit.')
        self.start = time.time()
        return self
 
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end = time.time()
        self.total = (self.end - self.start) * timeme.__unitfactor[self.unit]
        self.total = round(self.total, self.precision)
 
    def __str__(self):
        return 'Running time is {0}{1}'.format(self.total, self.unit)

实现了基本的时间单位设置、精确位数设置，并提供了默认的输出字符串。

from timeme import timeme
with timeme('ms', 6) as t:
    result = sum(range(100000))
print(r) # 4999950000
print(t.total) # 5.503178
print(t) # Running time is 5.503178ms

在一些简单的计时任务中，相比于使用timeit模块更加灵活而且无需手动配置运行环境；相比于使用手动计时减少了代码行数。

Teardown

• 前几天闲得无聊写了这个模块，很多都是得益于Python本身机制的便利。
• 还有很多值得改善的地方，例如重复测试timeit.repeat()的功能支持。
• 在代码代码计时、运行效率瓶颈检测方面，Python有很多完善的Profiler工具提供了相当强大的功能，我对这方面还不是很熟悉，最近打算学习一下。