Search code examples
pythonsqlalchemypython-decoratorssqlmodel

Python/SQLModel; nice syntax to add a method call to `__init__` and `update`?

I'm experimenting with SQLModel as an ORM. A few of my models have custom validation, calculated fields, or just things that I want to happen when they're created or changed. I end up using the following boilerplate a lot:

class MyModel(SqlModel):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.custom_method()

    def update(self, **kwargs):
        super().update(**kwargs)
        self.custom_method()

    def custom_method(self):
        """Do this when a model is created or updated
        """
        pass

Is there some nice way I can sweeten this syntax a little? Ideally I'd like a decorator around the function that would inject the function call into __init__ and update:

class MyModel(SqlModel):
    @run_on_change
    def custom_method(self):
        """Do this when a model is created or updated
        """
        pass

But I can't figure out how this would work, since a decorator intercepts when a function is called and modifies its behaviour, whereas I want to modify the circumstances in which the function is called.

Alternatively, can anyone make a compelling argument for using a @listens_for decorator instead of the above boilerplate approach associated with the model itself?

Solution

  • Define an abstract base class to act as the bridge between SQLModel and your concrete subclasses.

    from abc import ABC, abstractmethod


class MyBaseModel(SQLModel, ABC):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.custom_method()

    def update(self, **kwargs):
        super().update(**kwargs)
        self.custom_method()

    @abstractmethod
    def custom_method(self):
        pass


class MyModel(MyBaseModel):
    def custom_method(self):
        """Do stuff here"""