How to extend SQLalchemy Base class with a static method-CodePudding

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"
    ...