欢迎访问宙启技术站
智能推送

Python中的装饰器函数及其使用方式

发布时间:2023-06-20 13:11:37

装饰器函数是Python中一种特殊的函数,它可以在不改变函数本身的情况下,为函数添加附加的功能。装饰器在Python中被广泛使用,如Flask、Django等Web框架中的路由函数装饰器、Python内置的@staticmethod、@classmethod、@property等等。本文将详细介绍Python中的装饰器函数及其使用方式。

一、装饰器函数的定义

Python中的装饰器函数是一个可以接受一个函数作为参数,并且返回一个函数的函数。它通常用于修饰已有的函数,为其添加一些额外的功能,如输出函数执行的时间、记录日志等等。

下面是一个最简单的装饰器函数的定义:

def decorator(func):
    def wrapper(*args, **kwargs):
        print('Before function execution')
        func(*args, **kwargs)
        print('After function execution')
    return wrapper

其中,decorator函数接受一个函数作为参数,然后返回一个名为wrapper的函数。wrapper函数的参数和装饰器函数的参数一致,并且在函数执行前后分别输出一段信息。这个简单的装饰器已经可以用来修饰一个函数了。

如下面的例子,我们可以用上面的decorator函数来修饰一个名为test的函数:

@decorator
def test():
    print('Test')

这个语法表示将test函数作为参数传递给decorator函数,并将其返回的wrapper函数作为新的test函数。我们可以看一下新的test函数的输出:

>>> test()
Before function execution
Test
After function execution

这说明我们已经成功用decorator函数来修饰了test函数,为其添加了额外的功能。

二、装饰器函数的应用

以上我们已经定义了一个简单的装饰器函数,并用它来修饰了一个函数。下面我们将会讲到装饰器函数的应用。

1、记录函数执行时间

下面的装饰器函数可以用来记录一个函数的执行时间:

import time

def timeit(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(f'{func.__name__} function takes {end - start:.3f} seconds to run.')
        return result
    return wrapper

其中,timeit函数可以接受任意一个函数作为参数,返回一个新的包装函数wrapper。wrapper函数在执行调用原函数前后分别记录时间,并输出函数执行时间。下面我们利用这个装饰器函数来记录一个斐波那契数列函数的执行时间:

@timeit
def fib(n):
    if n <= 1:
        return n
    return fib(n - 1) + fib(n - 2)

print(fib(35))
# Output: fib function takes 3.277 seconds to run.

输出结果表明,fib函数执行了3.277秒。可以看到,我们使用了一个装饰器函数,就可以轻松地实现了记录函数执行时间的功能。

2、授权访问

下面的装饰器函数可以用来实现对某些需要授权才能访问的资源进行授权:

def authorize(func):
    def wrapper(*args, **kwargs):
        print('Check authorization')
        if check_auth():
            return func(*args, **kwargs)
        else:
            raise PermissionError('Unauthorized access')
    return wrapper

def check_auth():
    # Implementation of authorization check
    return True

其中,authorize函数取一个函数作为参数,然后返回一个新的包装函数wrapper。wrapper函数在执行调用原函数前先检查授权,如果授权成功则调用原函数,否则抛出权限错误。

下面我们利用这个装饰器函数来实现对一个资源的授权访问:

@authorize
def protected_resource():
    print('Access granted')

可以看到,protected_resource函数被装饰器authorize修饰,当调用该函数时,先检查授权是否有效,再执行实际的函数操作。

3、数据库连接池

下面的装饰器函数可以用来实现带有数据库连接池的资源:

import psycopg2
from contextlib import contextmanager

db_config = {
    'database': 'test_db',
    'user': 'test_user',
    'password': 'test_pwd',
    'host': 'localhost',
    'port': '5432'
}

@contextmanager
def db_cursor():
    conn = psycopg2.connect(**db_config)
    try:
        with conn:
            with conn.cursor() as cur:
                yield cur
    finally:
        conn.close()

def with_db_cursor(func):
    def wrapper(*args, **kwargs):
        with db_cursor() as cursor:
            kwargs['cursor'] = cursor
            return func(*args, **kwargs)
        
    return wrapper

其中,with_db_cursor装饰器函数连接到一个数据库,并返回一个新的包装函数wrapper。wrapper函数通过调用db_cursor函数来获取数据库连接,并将其作为参数传递给实际的函数。

下面我们定义一个利用数据库连接的实际函数,并使用with_db_cursor装饰器来实现自动获取和释放数据库连接的效果:

@with_db_cursor
def count_num_rows(cursor):
    cursor.execute('SELECT COUNT(*) FROM test_table')
    return cursor.fetchone()[0]

print(count_num_rows())

输出结果是test_table表中的行数,通过使用装饰器我们无需显式地获取和释放数据库连接,从而让代码更加简洁易读。

三、总结

在Python中,装饰器函数是一种非常有用的编程技巧,它可以让平常的函数具有更多的功能。本文主要介绍了Python中装饰器函数的定义、使用方法和应用场景。通过合理地使用装饰器函数,我们可以写出更加简洁、可读性更高的Python代码。