在不修改其代码的情况下向库函数添加装饰器

Question

我有一个第三方库函数，我的代码中很多地方都在使用它。现在我不想去所有地方都换衣服。所以我想猴子修补这个功能会帮助我实现我想要的装饰器。我只想在函数中添加一些具有重试机制的异常处理。

使用 Monkey Patching 的问题是递归。

按照我 project folder

正下方的文件夹结构

plugins
    __init__.py
    hooks
        __init__.py
        base_hook.py
        custom_ssh_hook.py
        ssh_hook.py

这就是我的代码的样子。

hooks/__init__.py

from plugins.hooks import ssh_hook
from plugins.hooks import custom_ssh_hook

#
print "running monkey patching ssh hook get_conn"
ssh_hook.SSHHook.get_conn = custom_ssh_hook.get_conn
print "completed monkey patching ssh hook get_conn"

hooks/base_hook.py

class BaseHook(object):
    def __init__(self, source):
        pass

hooks/ssh_hook.py

from plugins.hooks.base_hook import BaseHook

class SSHHook(BaseHook):
    def __init__(self, source, timeout=10):
        super(SSHHook, self).__init__(source)

    def get_conn(self):
        print("SSH Hook")

hooks/custom_ssh_hook.py

from plugins.hooks.ssh_hook import SSHHook

call_count = 0
def get_conn(self):
    global call_count
    call_count += 1
    if call_count > 1:
        print("Not A good Idea, you are trying recursion")
        return
    try:
        print("custom ssh Hook")
        return SSHHook.get_conn(self)
    except Exception as e:
        print("retry mechanism")
        raise e

我无法打印变量 originalHook 不管我如何导入它。我尝试关注

from plugins.hooks import originalHook
from . import *

get_conn of custom_ssh_hook 被递归调用。

如何调用 SSHHook 的 get_conn。

Answer 1

我找到了解决方案。为 class 我正在修补添加了一种方法。

猴子补丁

from airflow.contrib.hooks import ssh_hook
ssh_hook.SSHHook.original_func = ssh_hook.SSHHook.get_conn
from operators.bmo_ssh_hook import get_conn_with_retry


ssh_hook.SSHHook.get_conn = get_conn_with_retry

修补方法

max_retry = 3
retry = 0


def get_conn_with_retry(self):
    try:
        return SSHHook.original_func(self)
        # return self.original_func()
    except Exception as e:
        global retry
        if retry < max_retry:
            retry += 1
            print("tried %s times, failed to connect retrying again" %retry)
            self.get_conn()
        else:
            raise e

在不修改其代码的情况下向库函数添加装饰器

Adding decorator to library function without modifying its code

python

python-2.7