Source code for anyblok._graphviz

# This file is a part of the AnyBlok project
#
#    Copyright (C) 2014 Jean-Sebastien SUZANNE <jssuzanne@anybox.fr>
#    Copyright (C) 2022 Jean-Sebastien SUZANNE <js.suzanne@gmail.com>
#
# 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 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, }, )