How to decorate all callable methods of an imported lib

Question

I am using MLFlow lib to write data to a MLFlow instance running in the cloud.

Currently I have to deal with authentication before calling some methods of the lib 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(...)

The above code runs successfuly.

I would like that all callable mlflow methods first execute my custom auth logic, and only then execute the mlflow logic. For that I though of decorating all callable mlflow methods.

The idea is to do something like:

import mlflow
mlflow = authenticate_mlflow(mlflow)

mlflow.log_metric(...)

This way, when I run mlflow.log_metric(...), it would first execute my custom auth logic, since mlflow now is really the decorated mlflow, and only then it would make the mlflow API request behind the log_metric method.

To accomplish that I came up with the following code:

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(...)

But when I run the above code with python 3.11, it yields the error:

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

This is a good way of accomplishing the task "run my custom auth logic automatically before every callable mlflow method"? If so, what I am getting wrong in the implementation?

Answer 1

You've made the code unnecessarily complicated.

The logic is;

Get all attributes of the module.
Filter those attributes if that start with _ an underscore because they are private. Replace them if they're callable.

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))
    

Since the object passed to this function refers to a module object and we're setting it's attributes, there's no need to return the module from this function and reassign it to the module name.