本 日 の h o g e

■ [Py][小ネタ] SQLAlchemy でグラフ

相変わらず長いこと放置しているので，秋の夜長にと言うには少し遅いけど，何となく遊び．

ノードに種類を持たせつつ，多対多でノード同士を繋いでみたりするなど．

import uuid
import json

from contextlib import contextmanager

from sqlalchemy import create_engine
from sqlalchemy import Table, Column, ForeignKey, String
from sqlalchemy.orm import sessionmaker, relationship
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()
Engine = create_engine('sqlite:///:memory:')
Session = sessionmaker(bind=Engine)

NodeAssoc = Table(
    'node_assoc', Base.metadata,
    Column('parent_id', String(36),
           ForeignKey('node.id', onupdate='CASCADE', ondelete='CASCADE'),
           nullable=False),
    Column('child_id', String(36),
           ForeignKey('node.id', onupdate='CASCADE', ondelete='CASCADE'),
           nullable=False),
)


@contextmanager
def txn():
    session = Session()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


class ModelBase:
    def __init__(self, **kwargs):
        self.id = str(uuid.uuid4())
        self.apply(**kwargs)

    def apply(self, **kwargs):
        for key, val in kwargs.items():
            setattr(self, key, val)

    def dictify(self, ignore=['id']):
        return {n: getattr(self, n)
                for n in self.__table__.columns.keys() if n not in ignore}

    def __repr__(self):
        return json.dumps(self.dictify())


class Node(ModelBase, Base):
    TYPE = None

    id = Column(String(36), primary_key=True)
    type = Column(String(255), nullable=False)
    value = Column(String(255))
    children = relationship(
        'Node', secondary=NodeAssoc,
        primaryjoin=id == NodeAssoc.columns.parent_id,
        secondaryjoin=id == NodeAssoc.columns.child_id,
        backref='parents', uselist=True)

    __tablename__ = 'node'
    __mapper_args__ = {'polymorphic_on': type}

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.type = self.TYPE


class CircleNode(Node):
    TYPE = 'circle'
    __mapper_args__ = {'polymorphic_identity': TYPE}


class SquareNode(Node):
    TYPE = 'square'
    __mapper_args__ = {'polymorphic_identity': TYPE}


class TriangleNode(Node):
    TYPE = 'triangle'
    __mapper_args__ = {'polymorphic_identity': TYPE}


if __name__ == '__main__':
    def dump(mode, nodes):
        def _dump(mode, node, level=1):
            print(' ' * level, node)
            for n in getattr(node, mode):
                _dump(mode, n, level + 2)

        print(mode)
        for node in nodes:
            _dump(mode, node)
        print()

    Base.metadata.create_all(bind=Engine)

    with txn() as session:
        t1 = TriangleNode(value='1')
        t2 = TriangleNode(value='2')
        t3 = TriangleNode(value='3')
        s1 = SquareNode(value='1', children=[t1, t2])
        s2 = SquareNode(value='2', children=[t2, t3])
        c1 = CircleNode(value='1', children=[s1, s2])
        session.add(c1)

    with txn() as session:
        for mode, cls in [('children', CircleNode), ('parents', TriangleNode)]:
            dump(mode, session.query(cls).all())

実行するとこうなる．

children
  {"type": "circle", "value": "1"}
    {"type": "square", "value": "1"}
      {"type": "triangle", "value": "1"}
      {"type": "triangle", "value": "2"}
    {"type": "square", "value": "2"}
      {"type": "triangle", "value": "2"}
      {"type": "triangle", "value": "3"}

parents
  {"type": "triangle", "value": "1"}
    {"type": "square", "value": "1"}
      {"type": "circle", "value": "1"}
  {"type": "triangle", "value": "2"}
    {"type": "square", "value": "2"}
      {"type": "circle", "value": "1"}
    {"type": "square", "value": "1"}
      {"type": "circle", "value": "1"}
  {"type": "triangle", "value": "3"}
    {"type": "square", "value": "2"}
      {"type": "circle", "value": "1"}

だからどうしたぼくドラえもん．