How do I properly override `__getattr__` in SQLAlchemy to access dict keys as attributes?

Viewed 41

I'm trying to override __getattr__ in SQLAlchemy so I can access certain dictionary fields using object.field syntax rather than having to do object.fields["field"].

However, SQLAlchemy isn't recognizing changes to the field if I use the dot notation syntax. Accessing the field using normal object.fields["field"] however works as expected.

I've provided a code sample below.

#!/usr/bin/env python3
import sqlalchemy
from sqlalchemy import Column, Integer, JSON
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.ext.mutable import MutableDict
from sqlalchemy.orm import sessionmaker

Session = sessionmaker()
Base = declarative_base()


class Track(Base):  # noqa: WPS230
    __tablename__ = "track"

    id = Column(Integer, primary_key=True)
    _fields = Column(MutableDict.as_mutable(JSON), default="{}")

    def __init__(self, id):
        self.id = id
        self._fields = {}
        self._fields["custom"] = "field"

    def __getattr__(self, name: str):
        """See if ``name`` is a custom field."""
        try:
            return self.__dict__["_fields"][name]
        except KeyError:
            raise AttributeError from None

    def __setattr__(self, name, value):
        """Set custom fields if a valid key."""
        if "_fields" in self.__dict__ and name in self.__dict__["_fields"]:
            self._fields[name] = value
        else:
            super().__setattr__(name, value)


engine = sqlalchemy.create_engine("sqlite:///:memory:")
Session.configure(bind=engine)
Base.metadata.create_all(engine)  # creates tables
session = Session()

track = Track(id=1)
session.add(track)
session.commit()

assert track.custom
track.custom = "new" 
assert track in session.dirty
assert track.custom == "new"

I believe it has something to do with how sqlalchemy loads __dict__, but not sure how to work around it to get my desired behavior.

1 Answers

With the help of @CaseIIT on github, I came up with the following solution using a list FIELDS of fields that can be set:

class Track(Base):  # noqa: WPS230
    __tablename__ = "track"

    FIELDS = ["custom"]

    id = Column(Integer, primary_key=True)
    _fields = Column(MutableDict.as_mutable(JSON), default="{}")

    def __init__(self):
        self._fields = {}
        self._fields["custom"] = "field"

    def __getattr__(self, name: str):
        """See if ``name`` is a custom field."""
        if name in self.FIELDS:
            return self._fields[name]
        else:
            raise AttributeError from None

    def __setattr__(self, name, value):
        """Set custom custom_fields if a valid key."""
        if name in self.FIELDS:
            self._fields[name] = value
        else:
            super().__setattr__(name, value)

It feels wrong to have a duplicate list of accessible fields between _fields.keys() and FIELDS, but I can't seem to get around the issues that come with trying to inspect _fields.

Related