[Coding笔记] Python装饰器

官方说明：Python文档 - Decorator

A function returning another function, usually applied as a function transformation using the @wrapper syntax.

The decorator syntax is merely syntactic sugar, the following two function definitions are semantically equivalent:

def f(...):
    ...
f = staticmethod(f)

@staticmethod
def f(...):
    ...

一个简单的例子：

import functools


def log(func):
    @functools.wraps(func)  # 用于保留函数元信息
    def wrapper(*args, **kwargs):
        print('call %s():' % func.__name__)
        print('args = {}'.format(*args))
        return func(*args, **kwargs)

    return wrapper

调用时：

@log
def test(p):
    print(test.__name__ + " param: " + p)
    
test("I'm a param")

输出：
call test(): args = I'm a param test param: I'm a param

在实际项目中的用例：

from functools import wraps

def record_exception_related_info(pause=False):
    """
    记录发生异常时的相关信息，包括：1. 屏幕快照，2. 浏览器控制台日志
    ——注意：浏览器驱动必须作为测试函数所属的类/实例的属性，且命名为driver
    :param pause: 设为True时，程序将在发生异常时暂停，以便调试
    :return: 以装饰器的方式应用于测试函数
    """
    def exception_record_wrapper(test_func):
        @wraps(test_func)  # wraps用于保留函数的元信息
        def test_func_wrapper(*test_func_args, **test_func_kwargs):
            return ExceptionRecord(pause).call(test_func, *test_func_args, **test_func_kwargs)

        return test_func_wrapper

    return exception_record_wrapper


class ExceptionRecord:
    def __init__(self, pause):
        self.driver = None
        self.sleep_time = 10000 if pause else None

    def call(self, test_func, *test_func_args, **test_func_kwargs):
        self.driver = self.get_driver(test_func_args)
        try:
            test_func(*test_func_args, **test_func_kwargs)
        except:
            self.save_exception_related_information()
            if self.sleep_time:
                print(traceback.format_exc())
                time.sleep(self.sleep_time)
            raise

    def save_exception_related_information(self):
        current_date = datetime.datetime.strftime(datetime.datetime.now(), "%Y-%m-%d")
        current_time = datetime.datetime.strftime(datetime.datetime.now(), "%H_%M_%S")
        dir_path = f'{EXCEPTION_RECORD_DIR_PATH}\\{current_date}\\{current_time}'
        if not os.path.isdir(dir_path):
            os.makedirs(dir_path)
        self.driver.get_screenshot_as_file(f'{dir_path}\\screenshort.png')
        self.save_browser_console_data(dir_path)

    def save_browser_console_data(self, dir_path):
        for log_type in self.driver.log_types:
            log_data = pprint.pformat(self.driver.get_log(log_type))
            with open(f'{dir_path}\\{log_type}_log.txt', 'w') as file_pipeline:
                file_pipeline.write(log_data)

    @staticmethod
    def get_driver(test_func_args):
        # 参数检查，分别判断参数个数，参数名，参数类型是否正确
        if len(test_func_args) > 0 and DRIVER_ATTRIBUTE_NAME in dir(test_func_args[0]):
            driver = test_func_args[0].driver
            if isinstance(driver, WebDriver):
                return driver
            else:
                raise ValueError("测试类的driver属性不是浏览器驱动对象")
        else:
            raise ValueError("无法获得测试类的driver属性")


class _TestClassExample(TestCase):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 注意：浏览器驱动必须作为类属性，且命名为driver
        chrome_driver = ChromeDriver(headless=False, disable_w3c=True)
        self.driver = chrome_driver.driver

    @record_exception_related_info()
    def test_method_example(self):
        self.driver.get("http://www.baidu.com")
        self.driver.get("raise_exception")