Search code examples
pythonflaskflask-sqlalchemymarshmallow

How to handle Flask_SQLAlchemy models methods in Marshmallow schemas?


I have a few Flask_SQLAlchemy data models and their methods as follows:

class User(db.Model):

    __table_args__ = {'extend_existing': True}
    id = db.Column('id',db.Integer , primary_key=True)
    username = db.Column('username', db.String(), unique=False , index=True)
    password = db.Column('password' , db.String())
    email = db.Column('email',db.String(),unique=True , index=True)
    role = db.Column('role', db.Integer(), default=0)
    account_state = db.Column('account_state', db.Integer(), default=1)
    karma = db.relationship('Karma', backref='user', lazy='dynamic', cascade="all, delete-orphan")
    registered_on = db.Column('registered_on' , db.DateTime(timezone=True), server_default=func.now())

    followed = db.relationship('Follow', 
        foreign_keys=[Follow.follower_id], backref=db.backref('follower', lazy='joined'), 
        lazy='dynamic', cascade='all, delete-orphan')
    followers = db.relationship('Follow', foreign_keys=[Follow.followed_id], 
        backref=db.backref('followed', lazy='joined'), lazy='dynamic', cascade='all, delete-orphan')

    def to_dict(self):
        data = {
            'id': self.id,
            'username': self.username,
            'is_admin': self.is_admin(),
            'followers_count': self.followers_count(),
            'followed_count': self.followed_count(),
            'has_karma': self.has_karma(),
            'karma_count': self.karma_count(),
            'registered_on': self.registered_on,
        }
        return data

    def is_admin(self):
        if self.role == 1:
            return True
        return False

    def load_karma(self):
        karma = self.karma.filter_by(user_id=self.id).first()
        return karma

    def has_karma(self):
        karma = self.load_karma()
        if not karma:
            return False
        if karma:
            return True
        return False

    def karma_count(self):
        karma = self.load_karma()
        if karma:
            return karma.number
        return 0

    def followers_count(self):
        return self.followers.count()

    def followed_count(self):
        return self.followed.count()

class Karma(db.Model):
    __table_args__ = {'extend_existing': True}
    id = db.Column(db.Integer, primary_key=True)
    number = db.Column(db.Integer)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'))
    last_karma_date = db.Column(db.DateTime, default=datetime.utcnow)

class Follow(db.Model):
    __tablename__ = 'follows'
    __table_args__ = {'extend_existing': True}
    follower_id = db.Column(db.Integer, db.ForeignKey('user.id'), primary_key=True)
    followed_id = db.Column(db.Integer, db.ForeignKey('user.id'), primary_key=True)
    state = db.Column(db.Integer, default=0)
    date_followed = db.Column(db.DateTime, default=datetime.utcnow)

as you can see in the User class I have a few methods and one of which is to_dict() which I use to return JSON data for my API endpoints, but now that I switched to Marshmallow to automatically serailize my models, I could not understand how may I serialize the data my User model for example returns, as in has_karma() and karma_count(). I have been to the official docs and many online articles and questions here on stackoverflow but nothing yet. This is my User schema so far:

from marshmallow import Schema, fields, ValidationError

class UserSchema(Schema):
    id = fields.Int(dump_only=True)
    username = fields.Str()
    fullname = fields.Str()
    registered_on = fields.DateTime(dump_only=True)

How do I serialize the returned data of my methods as in to_dict()?


Solution

  • You have to change your method to an attribute. You can easily do this with the decorator @property.

    A minimal working example

    
    from flask import Flask
    from flask_sqlalchemy import SQLAlchemy
    from marshmallow import Schema, fields
    
    app = Flask(__name__)
    db = SQLAlchemy(app)
    
    
    class User(db.Model):    
        id = db.Column('id', db.Integer, primary_key=True)
        username = db.Column('username', db.String(), unique=False, index=True)
    
        @property
        def has_karma(self):
            return True
    
    
    class UserSchema(Schema):
        id = fields.Int(dump_only=True)
        username = fields.Str()
        has_karma = fields.Boolean()
    
    
    @app.route("/")
    def home():    
        return UserSchema().dump(User(id=1, username='hi'))
    
    
    if __name__ == '__main__':
        app.run()
    

    Output (if you run the app and visit http://127.0.0.1:5000/):

    has_karma   true
    id          1
    username    "hi"