Home > Back-end >  How to extend SQLalchemy Base class with a static method
How to extend SQLalchemy Base class with a static method

Time:12-29

I have multiple classes similar to the following:

class Weather(Base):
    __tablename__ = "Weather"
    id = Column(Integer, primary_key=True, nullable=False, autoincrement=True)
    temperature = Column(Integer)
    humidity = Column(Integer)
    wind_speed = Column(Float)
    wind_direction = Column(String)

I want to add a method df() that returns me the Pandas dataframe of that table. I know I can write it like this:

class Weather(Base):
   __tablename__ = "Weather"
    id = Column(Integer, primary_key=True, nullable=False, autoincrement=True)
    temperature = Column(Integer)
    humidity = Column(Integer)
    wind_speed = Column(Float)
    wind_direction = Column(String)

    @staticmethod
    def df():
        with engine.connect() as conn:
            return pd.read_sql_table(Weather.__tablename__ , conn)

But I want to implement this for every table. I guess if I can extend the Base class with this method I should be able to implement it once and use it in every class. Everything I have tried has failed because I do not have access to __tablename__ attribute.

SOLUTION I ended up with a mix of both answers. I have used the first method proposed by @snakecharmerb (it allows to introduce the change without modifying the rest of the code) with the @classmethod proposed by @RomanPerekhrest (which is the bit I was missing).

class MyBase:
    __tablename__ = None

    @classmethod
    def df(cls):
        with engine.connect() as conn:
            return pd.read_sql_table(cls.__tablename__ , conn)

Base = declarative_base(cls=MyBase)

CodePudding user response:

You can do this by passing a custom class to the declarative_base function:

class MyBase:
    __abstract__ = True

    @staticmethod
    def df():
        with engine.connect() as conn:
            return pd.read_sql_table(Weather.__tablename__ , conn)

Base = orm.declarative_base(cls=MyBase)

class MyModel(Base):
    __tablename__ = 'tbl'
    ...

Alternatively, you can create a mixin that provides the static method and have classes inherit from it selectively.


class DFMixin:
    @staticmethod
    def df():
        with engine.connect() as conn:
            return pd.read_sql_table(Weather.__tablename__ , conn)

class MyModel(Base, DFMixin):
    __tablename__ = 'tbl'
    ...

The mixin gives you more flexibility if not all of your models are going to need the dataframe functionality.

CodePudding user response:

Declare an auxiliary class (say DfBase) with classmethod df(cls) having the desired behavior.
Then each derived class will access its __tablename__ attribute seamlessly via cls object which refers to the derived class itself.

class DfBase:
    __tablename__ = None

    @classmethod
    def df(cls):
        with engine.connect() as conn:
            return pd.read_sql_table(cls.__tablename__ , conn)


class Weather(Base, DfBase):
    __tablename__ = "Weather"
    ...
  • Related