# This file is a part of the AnyBlok project
#
# Copyright (C) 2014 Jean-Sebastien SUZANNE <jssuzanne@anybox.fr>
#
# This Source Code Form is subject to the terms of the Mozilla Public License,
# v. 2.0. If a copy of the MPL was not distributed with this file,You can
# obtain one at http://mozilla.org/MPL/2.0/.
from graphviz.dot import Digraph
[docs]class BaseSchema:
"""Common class extended by the type of schema"""
def __init__(self, name, format='png'):
self.name = name
self.format = format
self._nodes = {}
self._edges = {}
self.count = 0
[docs] def add_edge(self, cls_1, cls_2, attr=None):
"""Add a new edge between two nodes
::
dot.add_edge(node1, node2)
:param cls_1: node (string or object) - from
:param cls_2: node (string or object) - to
:param attr: attribute of the edge
"""
cls_1 = cls_1 if isinstance(cls_1, str) else cls_1.name
cls_2 = cls_2 if isinstance(cls_2, str) else cls_2.name
self.count += 1
self._edges["%s_%s_2_%d" % (cls_1, cls_2, self.count)] = {
'from': cls_1,
'to': cls_2,
'attr': {} if attr is None else attr
}
[docs] def render(self):
"""Call graphviz to create the schema """
self.dot = Digraph(name=self.name, format=self.format,
node_attr={'shape': 'record',
'style': 'filled',
'fillcolor': 'gray95'})
for _, cls in self._nodes.items():
cls.render(self.dot)
for _, edge in self._edges.items():
self.dot.edge(edge['from'], edge['to'],
_attributes=edge['attr'])
[docs] def save(self):
"""Render and create the output file """
self.render()
self.dot.render(self.name)
[docs]class TableSchema:
"""Describe one table """
def __init__(self, name, parent, islabel=False):
self.name = name
self.parent = parent
self.islabel = islabel
self.column = []
[docs] def render(self, dot):
"""Call graphviz to create the schema """
if self.islabel:
label = "{%s}" % self.name
else:
column = '\\n'.join(self.column)
label = "{%s|%s}" % (self.name, column)
dot.node(self.name, label=label)
[docs] def add_column(self, name, type_, primary_key=False):
"""Add a new column to the table
:param name: the name of the column
:param type_: the type of the column
:param primary_key: if True, 'PK' argument will be added
"""
self.column.append("%s%s (%s)" % (
'PK ' if primary_key else '', name, type_))
[docs] def add_foreign_key(self, node, label=None, nullable=True):
"""Add a new foreign key
:param node: node (string or object) of the table attached
:param label: name of the column to add the foreign key to
TODO: i did not understand the explanation of 'nullable' parameter
:param nullable: boolean to select the multiplicity of the association
"""
self.parent.add_foreign_key(self, node, label, nullable)
[docs]class SQLSchema(BaseSchema):
"""Create a schema to display the table model
::
dot = SQLSchema('the name of my schema')
t1 = dot.add_table('Table 1')
t1.add_column('c1', 'Integer')
t1.add_column('c2', 'Integer')
t2 = dot.add_table('Table 2')
t2.add_column('c1', 'Integer')
t2.add_foreign_key(t1, 'c2')
dot.save()
"""
[docs] def add_table(self, name):
"""Add a new node TableSchema with columns
:param name: the name of the table
:rtype: returns an instance of TableSchema
"""
tmp = TableSchema(name, self)
self._nodes[name] = tmp
return tmp
[docs] def add_label(self, name):
"""Add a new node TableSchema without column
:param name: the name of the table
:rtype: returns an instance of TableSchema
"""
tmp = TableSchema(name, self, islabel=True)
self._nodes[name] = tmp
return tmp
[docs] def get_table(self, name):
"""Return the instance of TableSchema linked to the table name given
:param name: the name of the table
:rtype: return an instance of TableSchema
"""
return self._nodes.get(name)
def add_foreign_key(self, cls_1, cls_2, label=None, nullable=False):
multiplicity = "0..1" if nullable else "1"
hlabel = '%s (%s)' % (label, multiplicity) if label else multiplicity
self.add_edge(cls_1, cls_2, attr={
'arrowhead': "none",
'headlabel': hlabel,
})
[docs]class ClassSchema:
"""Used to display a class """
def __init__(self, name, parent, islabel=False):
self.name = name
self.parent = parent
self.islabel = islabel
self.properties = []
self.column = []
self.method = []
[docs] def extend(self, node):
"""Add an edge with extended shape to the node
:param node: node (string or object)
"""
self.parent.add_extend(self, node)
[docs] def strong_aggregate(self, node,
label_from=None, multiplicity_from=None,
label_to=None, multiplicity_to=None):
"""Add an edge with strong aggregate shape to the node
:param node: node (string or object)
:param label_from: the name of the attribute
:param multiplicity_from: multiplicity of the attribute
:param label_to: the name of the attribute
:param multiplicity_to: multiplicity of the attribute
"""
self.parent.add_strong_aggregation(self, node, label_from,
multiplicity_from, label_to,
multiplicity_to)
[docs] def aggregate(self, node,
label_from=None, multiplicity_from=None,
label_to=None, multiplicity_to=None):
"""Add an edge with aggregate shape to the node
:param node: node (string or object)
:param label_from: the name of the attribute
:param multiplicity_from: multiplicity of the attribute
:param label_to: the name of the attribute
:param multiplicity_to: multiplicity of the attribute
"""
self.parent.add_aggregation(self, node, label_from, multiplicity_from,
label_to, multiplicity_to)
[docs] def associate(self, node,
label_from=None, multiplicity_from=None,
label_to=None, multiplicity_to=None):
"""Add an edge with associate shape to the node
:param node: node (string or object)
:param label_from: the name of the attribute
:param multiplicity_from: multiplicity of the attribute
:param label_to: the name of the attribute
:param multiplicity_to: multiplicity of the attribute
"""
self.parent.add_association(self, node, label_from, multiplicity_from,
label_to, multiplicity_to)
[docs] def add_property(self, name):
"""Add a property to the class
:param name: the name of the property
"""
self.properties.append(name)
[docs] def add_column(self, name):
"""Add a column to the class
:param name: the name of the column
"""
self.column.append(name)
[docs] def add_method(self, name):
"""Add a method to the class
:param name: the name of the method
"""
self.method.append(name)
[docs] def render(self, dot):
"""Call graphviz to create the schema """
if self.islabel:
label = "{%s}" % self.name
else:
properties = '\\n'.join(self.properties)
column = '\\n'.join(self.column)
method = '\\n'.join('%s()' % x for x in self.method)
label = "{%s|%s|%s|%s}" % (self.name, properties, column, method)
dot.node(self.name, label=label)
[docs]class ModelSchema(BaseSchema):
"""Create a schema to display the UML model
::
dot = ModelSchema('The name of my UML schema')
cls = dot.add_class('My class')
cls.add_method('insert')
cls.add_property('items')
cls.add_column('my column')
dot.save()
"""
[docs] def add_class(self, name):
"""Add a new node ClassSchema with column
:param name: the name of the class
:rtype: return an instance of ClassSchema
"""
tmp = ClassSchema(name, self)
self._nodes[name] = tmp
return tmp
[docs] def add_label(self, name):
"""Return an instance of ClassSchema linked to the class name given
:param name: the name of the class
:rtype: return an instance of ClassSchema
"""
tmp = ClassSchema(name, self, islabel=True)
self._nodes[name] = tmp
return tmp
[docs] def get_class(self, name):
"""Add a new node ClassSchema without column
:param name: the name of the class
:rtype: return an instance of ClassSchema
"""
return self._nodes.get(name)
[docs] def add_extend(self, cls_1, cls_2):
"""Add edge to extend
:param cls_1: the name of the class 1
:param cls_2: the name of the class 2
"""
self.add_edge(cls_1, cls_2, attr={
'dir': 'back',
'arrowtail': 'empty',
})
[docs] def add_aggregation(self, cls_1, cls_2,
label_from=None, multiplicity_from=None,
label_to=None, multiplicity_to=None):
"""Add edge for aggregation
:param cls_1: the name of the class 1
:param cls_2: the name of the class 2
:param label_from: attribute name
:param multiplicity_from: multiplicity of the attribute
:param label_to: attribute name
:param multiplicity_to: multiplicity of the attribute
:return:
"""
label_from, label_to = self.format_label(
label_from, multiplicity_from, label_to, multiplicity_to)
if not cls_1 or not cls_2:
return # pragma: no cover
self.add_edge(cls_1, cls_2, attr={
'dir': 'back',
'arrowtail': 'odiamond',
'headlabel': label_from,
'taillabel': label_to,
})
[docs] def add_strong_aggregation(self, cls_1, cls_2,
label_from=None, multiplicity_from=None,
label_to=None, multiplicity_to=None):
"""Add edge for strong aggregation
:param cls_1:
:param cls_2:
:param label_from:
:param multiplicity_from:
:param label_to:
:param multiplicity_to:
:return:
"""
label_from, label_to = self.format_label(
label_from, multiplicity_from, label_to, multiplicity_to)
self.add_edge(cls_1, cls_2, attr={
'dir': 'back',
'arrowtail': 'diamond',
'headlabel': label_from,
'taillabel': label_to,
})
@staticmethod
def format_label(label_from, multiplicity_from, label_to,
multiplicity_to):
def _format_label(label, multiplicity):
if label:
if multiplicity:
return '%s (%s)' % (label, multiplicity)
return label
else:
if multiplicity:
return multiplicity
return
return (
_format_label(label_from, multiplicity_from),
_format_label(label_to, multiplicity_to),
)
def add_association(self, cls_1, cls_2,
label_from=None, multiplicity_from=None,
label_to=None, multiplicity_to=None):
label_from, label_to = self.format_label(
label_from, multiplicity_from, label_to, multiplicity_to)
self.add_edge(cls_1, cls_2, attr={
'arrowhead': "none",
'headlabel': label_from,
'taillabel': label_to,
})