Source code for anyblok._graphviz

# This file is a part of the AnyBlok project
#
#    Copyright (C) 2014 Jean-Sebastien SUZANNE <[email protected]>
#
# 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 new edge between 2 node :: dot.add_edge(node1, node2) :param cls_1: node (string or object) for the from :param cls_2: node (string or object) for the to :paam 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 do 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 in the table :param name: name of the column :param type_: type of the column :param primary_key: if True, the string PK will be add """ 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 linked :param label: name of the column of the foreign key :param nullable: bool 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 column :param name: name of the table :rtype: return the 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: name of the table :rtype: return the 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 with the name of table :param name: name of the table :rtype: return the 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: """ Use 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 extend shape to the node :param node: node (string or object) """ self.parent.add_extend(self, node)
[docs] def strong_agregate(self, node, label_from=None, multiplicity_from=None, label_to=None, multiplicity_to=None): """ add an edge with strong agregate shape to the node :param node: node (string or object) :param label_from: attribute name :param multiplicity_from: multiplicity of the attribute :param label_to: attribute name :param multiplicity_to: multiplicity of the attribute """ self.parent.add_strong_agregation(self, node, label_from, multiplicity_from, label_to, multiplicity_to)
[docs] def agregate(self, node, label_from=None, multiplicity_from=None, label_to=None, multiplicity_to=None): """ add an edge with agregate shape to the node :param node: node (string or object) :param label_from: attribute name :param multiplicity_from: multiplicity of the attribute :param label_to: attribute name :param multiplicity_to: multiplicity of the attribute """ self.parent.add_agregation(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: attribute name :param multiplicity_from: multiplicity of the attribute :param label_to: attribute name :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 in the class :param name: name of the property """ self.properties.append(name)
[docs] def add_column(self, name): """ add a column in the class :param name: name of the column """ self.column.append(name)
[docs] def add_method(self, name): """ add a method in the class :param name: name of the method """ self.method.append(name)
[docs] def render(self, dot): """Call graphviz to do 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: name of the class :rtype: return the instance of ClassSchema """ tmp = ClassSchema(name, self) self._nodes[name] = tmp return tmp
[docs] def add_label(self, name): """ Return the instance of ClassSchema linked with the name of class :param name: name of the class :rtype: return the 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: name of the class :rtype: return the instance of ClassSchema """ return self._nodes.get(name)
def add_extend(self, cls_1, cls_2): self.add_edge(cls_1, cls_2, attr={ 'dir': 'back', 'arrowtail': 'empty', }) def add_agregation(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) if not cls_1 or not cls_2: return self.add_edge(cls_1, cls_2, attr={ 'dir': 'back', 'arrowtail': 'odiamond', 'headlabel': label_from, 'taillabel': label_to, }) def add_strong_agregation(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={ 'dir': 'back', 'arrowtail': 'diamond', 'headlabel': label_from, 'taillabel': label_to, }) def format_label(self, 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, })