Python装饰器笔试题（简单难度）

前言

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

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

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

用函数调用函数

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

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

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)

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

用函数返回函数

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

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的时候又调用了一遍。

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

用函数返回包装函数

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

试试看：

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中访问到你的方法。

那么，现在我们应该如何使用呢？

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

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

看来@只是一个简写。

再进一步

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

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内容，然后再将这些内容写入文件。

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

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装饰器，这里就不再赘述了。这里实际上就是在双层嵌套的基础上再加一层嵌套，使得装饰器可以接收参数，就像这样：

@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字样了。