最近在看别人写的自动化代码,有好多重复的。比如很多操作需要login, logout。或者需要session等,看着这一段又一段重复的代码,心想是否可以写个装饰器来包装一下。
装饰器是什么?
python的装饰器本质上是一个Python函数,它可以让其他函数在不需要做任何代码变动的前提下增加额外功能,装饰器的返回值也是一个函数对象。简单的说装饰器就是一个用来返回函数的函数。
它经常用于有切面需求的场景,比如:插入日志、性能测试、事务处理、缓存、权限校验等场景。装饰器是解决这类问题的绝佳设计,有了装饰器,我们就可以抽离出大量与函数功能本身无关的雷同代码并继续重用。
概括的讲,装饰器的作用就是为已经存在的对象添加额外的功能。
在代码运行期间动态增加功能的方式,称之为“装饰器”(Decorator)。
可以这样理解,装饰器就是就是我们超市里买东西的购物袋,需要用它来装我们的东西。我们买了额外的东西,可以直接放到这个袋子里面,不需要改变以前的东西。
我们假设你的程序实现了say_hello()和say_goodbye()两个函数。
def say_hello():
print "hello!"
def say_goodbye():
print "hello!" # bug here
if __name__ == '__main__':
say_hello()
say_goodbye()
但是在实际调用中,我们发现程序出错了,上面的代码打印了两个hello。经过调试你发现是say_goodbye()出错了。老板要求调用每个方法前都要记录进入函数的名称,比如这样:
[DEBUG]: Enter say_hello()
Hello!
[DEBUG]: Enter say_goodbye()
Goodbye!
好,小A是个菜鸟,他是这样实现的。
def say_hello():
print "[DEBUG]: enter say_hello()"
print "hello!"
def say_goodbye():
print "[DEBUG]: enter say_goodbye()"
print "hello!"
if __name__ == '__main__':
say_hello()
say_goodbye()
很low吧? 嗯是的。这样就有大量重复的代码,影响到以前的函数了。 小A改进了一下。
def debug():
import inspect
caller_name = inspect.stack()[1][3]
print "[DEBUG]: enter {}()".format(caller_name)
def say_hello():
debug()
print "hello!"
def say_goodbye():
debug()
print "goodbye!"
if __name__ == '__main__':
say_hello()
say_goodbye()
是不是好一点?那当然,但是每个业务函数里都要调用一下debug()函数,是不是很难受?仍然影响到了以前的函数。
那么装饰器这时候应该登场了。
装饰器的作用就是为已经存在的函数或对象添加额外的功能。
def debug(func):
def wrapper():
print "[DEBUG]: enter {}()".format(func.__name__)
return func()
return wrapper
@debug
def say_hello():
print "hello!"
@debug
def say_goodbye():
print "goodbye!"
上面的debug函数其实已经是一个装饰器了,相当是我们套了两层包装袋,里面一层(wrapper())是用来包装别的东西,外面一层(debug(func))是来总的包装,当然包装袋的名字是可以自己取的。每层都要返回当层的函数名。
装饰器语法糖
python提供了@符号作为装饰器的语法糖,使我们更方便的应用装饰函数。但使用语法糖要求装饰函数必须return一个函数对象。因此我们将上面的func函数使用内嵌函数包裹并return。
装饰器相当于执行了装饰函数debug后又返回被装饰函数say_hello,因此say_hello()被调用的时候相当于执行了两个函数。
装饰器参数
这是最简单的装饰器,但是有一个问题,如果被装饰的函数需要传入参数,那么这个装饰器就坏了。因为返回的函数并不能接受参数,你可以指定装饰器函数wrapper接受和原函数一样的参数,比如:
def debug(func):
def wrapper(something): # 指定一毛一样的参数
print "[DEBUG]: enter {}()".format(func.__name__)
return func(something)
return wrapper # 返回包装过函数
@debug
def say(something):
print "hello {}!".format(something)
这样你就解决了一个问题,但又多了N个问题。因为函数有千千万,你只管你自己的函数,别人的函数参数是什么样子,鬼知道?还好Python提供了可变参数args和关键字参数*kwargs,有了这两个参数,装饰器就可以用于任意目标函数了。
def debug(func):
def wrapper(*args, **kwargs): # 指定宇宙无敌参数
print "[DEBUG]: enter {}()".format(func.__name__)
print 'Prepare and say...',
return func(*args, **kwargs)
return wrapper # 返回
@debug
def say(something):
print "hello {}!".format(something)
至此,你已完全掌握初级的装饰器写法。
带参数的装饰器
假设我们前文的装饰器需要完成的功能不仅仅是能在进入某个函数后打出log信息,而且还需指定log的级别,那么装饰器就会是这样的。
def logging(level):
def wrapper(func):
def inner_wrapper(*args, **kwargs):
print "[{level}]: enter function {func}()".format(
level=level,
func=func.__name__)
return func(*args, **kwargs)
return inner_wrapper
return wrapper
@logging(level='INFO')
def say(something):
print "say {}!".format(something)
# 如果没有使用@语法,等同于
# say = logging(level='INFO')(say)
@logging(level='DEBUG')
def do(something):
print "do {}...".format(something)
if __name__ == '__main__':
say('hello')
do("my work")
是不是有一些晕?你可以这么理解,当带参数的装饰器被打在某个函数上时,比如@logging(level=’DEBUG’),它其实是一个函数,会马上被执行,只要这个它返回的结果是一个装饰器时,那就没问题。细细再体会一下。
如果开始被包装的函数有参数,跟包装袋没有关系,你让它自己传一个自己的万能参数就可了。如果要对装饰器加参数,相当是又买了东西,需要多加一层包装袋。
基于类实现的装饰器
装饰器函数其实是这样一个接口约束,它必须接受一个callable对象作为参数,然后返回一个callable对象。在Python中一般callable对象都是函数,但也有例外。只要某个对象重载了call()方法,那么这个对象就是callable的。
class Test():
def __call__(self):
print 'call me!'
t = Test()
t() # call me
像call这样前后都带下划线的方法在Python中被称为内置方法,有时候也被称为魔法方法。重载这些魔法方法一般会改变对象的内部行为。上面这个例子就让一个类对象拥有了被调用的行为。
回到装饰器上的概念上来,装饰器要求接受一个callable对象,并返回一个callable对象(不太严谨,详见后文)。那么用类来实现也是也可以的。我们可以让类的构造函数init()接受一个函数,然后重载call()并返回一个函数,也可以达到装饰器函数的效果。
class logging(object):
def __init__(self, func):
self.func = func
def __call__(self, *args, **kwargs):
print "[DEBUG]: enter function {func}()".format(
func=self.func.__name__)
return self.func(*args, **kwargs)
@logging
def say(something):
print "say {}!".format(something)
带参数的类装饰器
如果需要通过类形式实现带参数的装饰器,那么会比前面的例子稍微复杂一点。那么在构造函数里接受的就不是一个函数,而是传入的参数。通过类把这些参数保存起来。然后在重载call方法是就需要接受一个函数并返回一个函数。
class logging(object):
def __init__(self, level='INFO'):
self.level = level
def __call__(self, func): # 接受函数
def wrapper(*args, **kwargs):
print "[{level}]: enter function {func}()".format(
level=self.level,
func=func.__name__)
func(*args, **kwargs)
return wrapper #返回函数
@logging(level='INFO')
def say(something):
print "say {}!".format(something)
装饰器执行顺序
多个装饰器执行的顺序就是从最后一个装饰器开始,执行到第一个装饰器,再执行函数本身。
def dec1(func):
print("1111")
def one():
print("2222")
func()
print("3333")
return one
def dec2(func):
print("aaaa")
def two():
print("bbbb")
func()
print("cccc")
return two
@dec1
@dec2
def test():
print("test test")
test()
执行结果:
aaaa
1111
2222
bbbb
test test
cccc
3333
写了这么多,我们可以看一个简单的例子。例如selenium失败截图。
# coding:utf-8
from selenium import webdriver
driver = webdriver.Firefox()
# 截图功能
def get_screen():
'''截图'''
import time
nowTime = time.strftime("%Y_%m_%d_%H_%M_%S")
driver.get_screenshot_as_file('%s.jpg' % nowTime)
# 自动截图装饰器
def screen(func):
'''截图装饰器'''
def inner(*args, **kwargs):
try:
f = func(*args, **kwargs)
return f
except:
get_screen() # 失败后截图
return inner
@screen
def search(driver):
driver.get("https://www.baidu.com")
driver.find_element_by_id("kw11").send_keys("python") # 此行运行失败的
driver.find_element_by_id("su").click()
search(driver) # 执行search
我们不需要每个地方都去调用截图函数,只是在需要的地方,加上@screen,函数如果失败,就会自动截图。比我们手动添加只能且代码没有那么多冗余。
更多精彩,请关注微信公众号:python爱好部落