Pickle dynamically generated classes with flexible parent classes?

Question

There are a family of classes for which I wish to optionally extend and override some functionality. For these classes, I wish to add the same functionality, as they are all compatible. I am using Python 3.8+. The way I achieved this was by creating the class as a type with the additional functionality and passing the parent class as the bases. As a basic example:

class A:
    def __init__(self, a, **kwargs):
        self.a = a
        self.params = kwargs

class B:
    def __init__(self, b, **kwargs):
        self.b = b
        self.params = kwargs

def extend_class_with_print_params(base_class):
    def print_params(self):
        for k, v in self.params.items():
            print(k, v)

    return type(
        f"Extended{base_class.__name__}",
        (base_class,),
        dict(print_params=print_params),
    )

In the above, I define A and B. The function extend_class_with_print_params adds functionality compatible with both. My actual use case is adding pre-train and post-predict hooks to some instances of specific sklearn predictors, which is why I need the parent to be configurable.

import joblib
from test_classes import *

normal_a = A(a=10)
joblib.dump(normal_a, "normal_a")
normal_a = joblib.load("normal_a")

extended_a = extend_class_with_print_params(A)(a=15, alpha=0.1, beta=0.2)
joblib.dump(extended_a, "extended_a")
extended_a = joblib.load("extended_a")

When dumping extended_a, the following error is thrown:

_pickle.PicklingError: Can't pickle <class 'test_classes.ExtendedA'>: it's not found as test_classes.ExtendedA

As suggested in one of the below posts, I attempted setting new_class_name in globals to point to the new class before returning in the function. This allowed me to successfully dump, but not load the file in a different session, which makes sense since the globals would be reset. In general, I would also prefer not to modify globals anyway.

I have tried but failed to work out a solution using __reduce__ based on the following:

• Pickling dynamically generated classes?
• How can I pickle a dynamically created nested class in python?
• Pickle a dynamically parameterized sub-class

I didn't find the above methods clearly to apply to my situation. The content may be relevant and directly applicable, but I failed to find a way.

I'm also entirely open to changing my pattern (even if it means not dynamically defining the class). In short, I have the following requirements:

• Extend and override some arbitrary parent class's functionality, but not in all cases, since it will be optional whether to extend/override the class
• The objects must be pickle-able
• The objects must be pickled using joblib or the pickle library, not something like cloudpickle

Answer 1

It's probably best to avoid dynamically generating a class. Ideally, you can account for added functionality from the beginning. If you have control over classes A and B, you could do a pattern like this:

class A:
    hook: Callable
    def __init__(self, b, **kwargs):
        self.b = b
        self.params = kwargs

    def print_param_hook(self):
        if self.hook:
            self.hook(self.params.items())
        else:
            raise ArithmeticError("No hook function supplied!")

    def set_hook(self, hook: Callable):
        self.hook = hook

    def hook(items):
        for k, v in items:
            print(k, v)

a = A("foo", y="bar")
a.set_hook(hook1)

a.print_param_hook()

Here, A is defined with a pre-existing method that will call a generic function provided by the user. This, of course, constrains what sort of arguments your hook function can take.

Another option is to make a subclass of A and add your method to the subclass. Continuing the above example:

class SubA(A):
    def print_params(self):
        for k, v in self.params.items():
            print(k, v)

subA = SubA("foo", y="bar")
subA.print_params()

Finally, if you must add an arbitrary method to a class, you can do this using setattr:

def attr_hook(self):
    for k, v in self.params.items():
        print(k, v)  
setattr(A, 'attr_hook', attr_hook)
new_a = A("foo", y="bar")
new_a.attr_hook()

Note that this will affect every instance of A created, including those created before setattr, which isn't super desirable. You can read more about using setattr in this way in this blog post, including how to make a decorator to make it more seamless.

All of the options are completely pickleable:

import pickle

with open("test.pyc", "wb") as file:
    pickle.dump(new_a, file)

with open("test.pyc", "rb") as file:
    b = pickle.load(file)
b.attr_hook()