如何装饰导入库的所有可调用方法

Question

我正在使用MLFlow库，将数据写入一个在云端运行的MLFlow实例。

目前，在调用这个库的某些方法之前，我需要先进行身份验证：

# First I deal with auth explicitly
os.environ["MLFLOW_TRACKING_TOKEN"] = my_custom_get_token_function()
...

# Only then I use the mlflow methods
mlflow.log_metric(...)

上面的代码运行得很成功。

我希望所有可以调用的mlflow方法都能先执行我自定义的身份验证逻辑，然后再执行MLFlow的逻辑。为此，我想到了给所有可以调用的mlflow方法加上装饰器。

我的想法是这样做：

import mlflow
mlflow = authenticate_mlflow(mlflow)

mlflow.log_metric(...)

这样，当我运行mlflow.log_metric(...)时，它会先执行我的自定义身份验证逻辑，因为此时的mlflow实际上是被装饰过的版本，然后才会执行log_metric方法背后的mlflow API请求。

为了实现这个目标，我想出了以下代码：

def authenticate_mlflow(cls):
    class AuthenticatedMlflow(cls):
        @staticmethod
        def __getattribute__(cls, name):
            attr = super().__getattribute__(name)
            if callable(attr):
                return authenticate_mlflow_method(attr)
            return attr

    return AuthenticatedMlflow

def authenticate_mlflow_method(func):
    def wrapper(*args, **kwargs):
        handle_auth() # this functions has my custom auth logic
        return func(*args, **kwargs)

    return wrapper

import mlflow
mlflow = authenticate_mlflow(mlflow)

mlflow.log_metric(...)

但是当我用Python 3.11运行上面的代码时，出现了错误：

    mlflow = authenticate_mlflow(mlflow)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
File ..., in authenticate_mlflow
    class AuthenticatedMlflow(cls):
...
TypeError: module() takes at most 2 arguments (3 given)

这样做是否是一个好方法，可以在每个可调用的mlflow方法之前自动运行我的自定义身份验证逻辑？如果是的话，我在实现上哪里出了问题呢？

Answer 1

你的代码写得有点复杂，其实可以简单一些。

逻辑是这样的：

首先，获取这个模块的所有属性。
然后，筛选出那些以_开头的属性，因为它们是私有的。
最后，如果这些属性是可以调用的，就替换掉它们。

def authenticate_mlflow(mod):
    for attr_name in dir(mod):
        if attr_name.startswith('_'):
            continue
        attr = getattr(mod, attr_name)
        if callable(attr):
            setattr(mod, attr_name, authenticate_mlflow_method(attr))
    

因为传给这个函数的对象是一个模块对象，而我们正在设置它的属性，所以其实没必要从这个函数返回模块，然后再把它重新赋值给模块的名字。