Python装饰器笔试题(简单难度)

前言

这次遇到了一个比较神奇的面试题:给定方法

1
2
def add(x, y):
return x + y

要求在不改变源代码的前提下,使用装饰器,为add方法增加运行时间输出的功能。

用函数调用函数

其实本身而言并没有什么特别难的内容,只是单纯的比较综合罢了。

首先,我们尝试一下函数包函数,也就是说,我们将add函数包装为其他的函数:

1
2
3
4
5
6
7
8
9
import time
def decorate_add(x, y):
start = time.time()
def add(x, y):
return x + y
result = add(x, y)
end = time.time()
time_consumption = end - start
print(result, time_consumption)

看起来没什么问题,连输出也不需要说明,一目了然。但是呢,这样会不会太简单了?我们先升级一下难度。

用函数返回函数

没错,函数返回函数,然后调用函数,就能够获得结果。就像这样:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import time

def time_consume_f(f):
start = time.time()
print(f())
end = time.time()
print(end - start)
return time_consume_f

def add():
return 1e3

decorater = time_consume_f(add)
decorater(add)

输出结果就会是

1000.0

0.0

1000.0

0.0

没错,因为f本身具有print功能,在调用decorater的时候又调用了一遍。

当然,你也看到,这样的方法没办法继续输入参数了。得再想想办法。

用函数返回包装函数

既然我们单纯的用一层函数包装不够,我们使用两层呢?外层先传入函数,内层使用相同的参数,并调用外层传入的函数,是不是就能解决问题呢?

试试看:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import time
def decorater(func):
def wrapped(*args, **kwargs):
start = time.time()
result = func(*args, **kwargs)
end = time.time()
print('time: ', end - start)
return result
return wrapped

def add(x, y):
return x + y

wrapped_function = decorater(add)
wrapped_function(1, 2)

输出:

0.0

3

这下是我们想要的结果了。

使用@

如果你使用Flask,你会想到你曾经在一些controller方法上增加一个@app.route('/api'),让前端从http://localhost:8080/api中访问到你的方法。

那么,现在我们应该如何使用呢?

其实并不需要做出太多改动:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import time
def decorater(func):
def wrapped(*args, **kwargs):
start = time.time()
result = func(*args, **kwargs)
end = time.time()
print('time: ', end - start)
return result
return wrapped

@decorater
def add(x, y):
return x + y

add(1, 2)

输出还是老样子:

0.0

3

看来@只是一个简写。

再进一步

其实到此为止已经足够了。但是既然官方都有,那我们就用一下吧。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import time
from functools import wraps

def time_consume(f):
@wraps(f)
def wrapTheFunction(*args, **kargs):
start = time.time()
result = f(*args, **kargs)
end = time.time()
print(end - start)
return result
return wrapTheFunction

@time_consume
def add(x, y):
return x + y

print(add(1, 2))

其实效果是完全相同的:

0.0

3

但是呢,@wraps的作用主要就是装饰一个函数,并赋值函数名称、注释文档、参数列表等等

接受一个函数来进行装饰,并加入了复制函数名称、注释文档、参数列表等等的功能。总之就是,比我们单纯用双层函数功能全多了。

应用场景

当然,我们现在只考虑了两层嵌套的用法,也就是对一个函数进行包装,从而无侵入地使得这个函数拥有更多的功能。

而如果我再嵌套一层呢?没错,还能继续接收参数。

这也就是Flask@app.route()中能够有那么多参数的真正原因。我们通过这样的三层嵌套实现相对较为统一的、相对比较重复的内容,从而让我们在实现业务的过程中能够更专注于业务,剩下的装饰一下就好了。

我们举个例子吧。比如,我需要将日志保存在指定的文件中,但是我又不希望在代码中植入日志编辑的逻辑,因为这样只会让代码更为复杂。那么该怎么办呢?当然就是使用装饰器,通过截取函数中的所有print内容,然后再将这些内容写入文件。

我们可以首先定义一个三层嵌套的装饰器:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import sys
from contextlib import redirect_stdout
from io import StringIO

def redirect_print_to_log(log_path="app.log"):
def decorator(func):
def wrapper(*args, **kwargs):
# 创建一个字符串流用于捕获print输出
temp_stdout = StringIO()
# 保存原始的sys.stdout
original_stdout = sys.stdout
try:
# 重定向标准输出到字符串流
with redirect_stdout(temp_stdout):
result = func(*args, **kwargs)
# 将捕获的输出追加到指定的日志文件中
with open(log_path, "w", encoding='utf-8') as log_file:
log_file.write(temp_stdout.getvalue())
finally:
# 恢复原始的sys.stdout
sys.stdout = original_stdout
return result
return wrapper
return decorator

当然,你也可以使用@wraps装饰器,这里就不再赘述了。这里实际上就是在双层嵌套的基础上再加一层嵌套,使得装饰器可以接收参数,就像这样:

1
2
3
@redirect_print_to_log("your_file.log")
def your_function(*args, **kwargs):
print("your print")

在这段代码中,装饰器接收的参数就是log_path,即日志文件的路径your_file.log。这样就能够截取your_function中所有的print输出(这里就是截取到了your print),然后将这些输出写入到your_file.log中。

最后,打开your_file.log,文件中就存在your print字样了。