ThePianoDentist/output.py

from redbaron import RedBaron

code = """ 

class Rank(db.Model):
    __tablename__ = 'rank'
    id = db.Column(db.Integer, primary_key=True, autoincrement=True)
    player = db.Column(db.Integer, db.ForeignKey(Player), nullable=False, index=True)
    region = db.Column(db.String(3), nullable=False, index=True)
    mode = db.Column(db.String(5), nullable=False, index=True)
    timestamp = db.Column(db.Timestamp, nullable=False, default=func.now())
    score_rank = db.Column(db.Integer, nullable=False)
    kill_rank = db.Column(db.Integer, nullable=False)
    win_rank = db.Column(db.Integer, nullable=False)
"""

rb = RedBaron(code)

for class_ in rb.find_all("class"):
    if class_.find("def", name="__init__"):
        continue

    arguments = ["self"]

    for column in class_.find_all("assignment"):
        col_name = column.name.dumps()
        if col_name[:2] == "__" or "default=" in column.value.dumps():
            continue

        arguments.append(col_name)

    class_.append("\n")
    class_.append("def __init__(self):\n    stuff\n")

    foo = class_.find("def", name="__init__")
    foo.arguments = ", ".join(arguments)

    for arg_ in arguments[1:]:  # dont add self to body
        foo.append("    self.{0} = {0}\n".format(arg_))

    print class_

class Rank(db.Model):
    __tablename__ = 'rank'
    id = db.Column(db.Integer, primary_key=True, autoincrement=True)
    player = db.Column(db.Integer, db.ForeignKey(Player), nullable=False, index=True)
    region = db.Column(db.String(3), nullable=False, index=True)
    mode = db.Column(db.String(5), nullable=False, index=True)
    timestamp = db.Column(db.Timestamp, nullable=False, default=func.now())
    score_rank = db.Column(db.Integer, nullable=False)
    kill_rank = db.Column(db.Integer, nullable=False)
    win_rank = db.Column(db.Integer, nullable=False)

    def __init__(self, id, player, region, mode, score_rank, kill_rank, win_rank):
        stuff
        self.id = id
        self.player = player
        self.region = region
        self.mode = mode
        self.score_rank = score_rank
        self.kill_rank = kill_rank
        self.win_rank = win_rank