[PATCH 16 of 24 yams V2] [mypy] type schema2dot.py

Laurent Peuch cortex at worlddomination.be
Wed Mar 4 15:17:52 CET 2020


# HG changeset patch
# User Laurent Peuch <cortex at worlddomination.be>
# Date 1580263945 -3600
#      Wed Jan 29 03:12:25 2020 +0100
# Node ID 7c2414f5ab7e8c1ee2ab5bb99de8b74c79b8e343
# Parent  e17b76d216a4c7cfe5eb73a9ecf6543ff9b890c1
# Available At https://hg.logilab.org/users/lpeuch/yams
#              hg pull https://hg.logilab.org/users/lpeuch/yams -r 7c2414f5ab7e
# EXP-Topic type_annotations
[mypy] type schema2dot.py

diff --git a/yams/schema2dot.py b/yams/schema2dot.py
--- a/yams/schema2dot.py
+++ b/yams/schema2dot.py
@@ -19,38 +19,49 @@
 
 """
 from __future__ import print_function
-__docformat__ = "restructuredtext en"
 
 import sys
 import os.path as osp
 from itertools import cycle
 
+from typing import Any, Dict, Generator, Tuple, TypeVar, Set, Sequence, Optional
+
 from logilab.common.graph import DotBackend, GraphGenerator
 
-CARD_MAP = {'?': '0..1',
-            '1': '1',
-            '*': '0..n',
-            '+': '1..n'}
+from yams.types import _EntitySchemaType, _RelationSchemaType, _SchemaType
+
+__docformat__: str = "restructuredtext en"
+
+_T = TypeVar('_T')
+
+CARD_MAP: Dict[str, str] = {'?': '0..1',
+                            '1': '1',
+                            '*': '0..n',
+                            '+': '1..n'}
 
 
 class SchemaDotPropsHandler(object):
-    def __init__(self, visitor):
+    def __init__(self, visitor) -> None:
         self.visitor = visitor
         # FIXME: colors are arbitrary
         self._colors = cycle(('#aa0000', '#00aa00', '#0000aa',
                               '#000000', '#888888'))
         self.nextcolor = lambda: next(self._colors)
 
-    def node_properties(self, eschema):
+    def node_properties(self, eschema: _EntitySchemaType) -> Dict[str, str]:
         """return default DOT drawing options for an entity schema"""
         label = ['{', eschema.type, '|']
+
         label.append(r'\l'.join(rel.type for rel in eschema.ordered_relations()
                                 if rel.final and self.visitor.should_display_attr(eschema, rel)))
+
         label.append(r'\l}')  # trailing \l ensure alignement of the last one
+
         return {'label': ''.join(label), 'shape': "record",
                 'fontname': "Courier", 'style': "filled"}
 
-    def edge_properties(self, rschema, subjnode, objnode):
+    def edge_properties(self, rschema: Optional[_RelationSchemaType],
+                        subjnode, objnode) -> Dict[str, str]:
         """return default DOT drawing options for a relation schema"""
         # rschema can be none if the subject is a specialization of the object
         # we get there because we want to draw a specialization arrow anyway
@@ -67,6 +78,7 @@ class SchemaDotPropsHandler(object):
             kwargs = {'label': rschema.type,
                       'color': 'black', 'style': 'filled'}
             rdef = rschema.rdef(subjnode, objnode)
+
             if rdef.composite == 'subject':
                 kwargs['arrowhead'] = 'none'
                 kwargs['arrowtail'] = 'diamond'
@@ -76,135 +88,168 @@ class SchemaDotPropsHandler(object):
             else:
                 kwargs['arrowhead'] = 'normal'
                 kwargs['arrowtail'] = 'none'
+
             # UML like cardinalities notation, omitting 1..1
             if rdef.cardinality[1] != '1':
                 kwargs['taillabel'] = CARD_MAP[rdef.cardinality[1]]
+
             if rdef.cardinality[0] != '1':
                 kwargs['headlabel'] = CARD_MAP[rdef.cardinality[0]]
+
             kwargs['color'] = self.nextcolor()
+
         kwargs['fontcolor'] = kwargs['color']
+
         # dot label decoration is just awful (1 line underlining the label
         # + 1 line going to the closest edge spline point)
         kwargs['decorate'] = 'false'
         # kwargs['labelfloat'] = 'true'
+
         return kwargs
 
 
 class SchemaVisitor(object):
-    def __init__(self, skiptypes=()):
-        self._done = set()
-        self.skiptypes = skiptypes
-        self._nodes = None
-        self._edges = None
+    def __init__(self, skiptypes: Sequence = ()) -> None:
+        self._done: Set[Tuple[Any, Any, Any]] = set()
+        self.skiptypes: Sequence = skiptypes
+        self._nodes: Optional[Set[Tuple[Any, Any]]] = None
+        self._edges: Optional[Set[Tuple[Any, Any, Any]]] = None
 
-    def should_display_schema(self, erschema):
+    def should_display_schema(self, erschema) -> bool:
         return not (erschema.final or erschema in self.skiptypes)
 
-    def should_display_attr(self, eschema, rschema):
+    def should_display_attr(self,
+                            eschema: _EntitySchemaType,
+                            rschema: _RelationSchemaType) -> bool:
         return rschema not in self.skiptypes
 
-    def display_rel(self, rschema, setype, tetype):
+    def display_rel(self, rschema: _RelationSchemaType, setype, tetype) -> bool:
         if (rschema, setype, tetype) in self._done:
             return False
+
         self._done.add((rschema, setype, tetype))
+
         if rschema.symmetric:
             self._done.add((rschema, tetype, setype))
+
         return True
 
-    def nodes(self):
+    def nodes(self) -> Any:
         return self._nodes
 
-    def edges(self):
+    def edges(self) -> Any:
         return self._edges
 
 
 class FullSchemaVisitor(SchemaVisitor):
-    def __init__(self, schema, skiptypes=()):
+    def __init__(self, schema: _SchemaType, skiptypes: Sequence = ()) -> None:
         super(FullSchemaVisitor, self).__init__(skiptypes)
         self.schema = schema
-        self._eindex = {eschema.type: eschema for eschema in schema.entities()
-                        if self.should_display_schema(eschema)}
+        self._eindex: Dict[str, Any] = {eschema.type: eschema for eschema in schema.entities()
+                                        if self.should_display_schema(eschema)}
 
-    def nodes(self):
+    def nodes(self) -> Generator[Tuple[Any, Any], Any, None]:
         for eschema in sorted(self._eindex.values(), key=lambda x: x.type):
             yield eschema.type, eschema
 
-    def edges(self):
+    def edges(self) -> Generator[Tuple[str, str, Any], Any, None]:
         # Entities with inheritance relations.
         for eschema in sorted(self._eindex.values(), key=lambda x: x.type):
             if eschema.specializes():
                 yield str(eschema), str(eschema.specializes()), None
+
         # Subject/object relations.
         for rschema in sorted(self.schema.relations(), key=lambda x: x.type):
             if not self.should_display_schema(rschema):
                 continue
+
             for setype, tetype in sorted(rschema.rdefs, key=lambda x: (x[0].type, x[1].type)):
                 if not (setype in self._eindex and tetype in self._eindex):
                     continue
+
                 if not self.display_rel(rschema, setype, tetype):
                     continue
+
                 yield str(setype), str(tetype), rschema
 
 
 class OneHopESchemaVisitor(SchemaVisitor):
-    def __init__(self, eschema, skiptypes=()):
+    def __init__(self, eschema: _EntitySchemaType, skiptypes: Sequence = ()) -> None:
         super(OneHopESchemaVisitor, self).__init__(skiptypes)
         nodes = set()
         edges = set()
+
         nodes.add((eschema.type, eschema))
+
         for rschema in eschema.subject_relations():
             if not self.should_display_schema(rschema):
                 continue
+
             for teschema in rschema.objects(eschema.type):
                 nodes.add((teschema.type, teschema))
                 if not self.display_rel(rschema, eschema.type, teschema.type):
                     continue
+
                 edges.add((eschema.type, teschema.type, rschema))
+
         for rschema in eschema.object_relations():
             if not self.should_display_schema(rschema):
                 continue
+
             for teschema in rschema.subjects(eschema.type):
                 nodes.add((teschema.type, teschema))
+
                 if not self.display_rel(rschema, teschema.type, eschema.type):
                     continue
+
                 edges.add((teschema.type, eschema.type, rschema))
+
         # Inheritance relations.
         if eschema.specializes():
             nodes.add((eschema.specializes().type, eschema.specializes()))
             edges.add((eschema.type, eschema.specializes().type, None))
+
         if eschema.specialized_by():
             for pschema in eschema.specialized_by():
                 nodes.add((pschema.type, pschema))
                 edges.add((pschema.type, eschema.type, None))
+
         self._nodes = nodes
         self._edges = edges
 
 
 class OneHopRSchemaVisitor(SchemaVisitor):
-    def __init__(self, rschema, skiptypes=()):
+    def __init__(self, rschema: _RelationSchemaType, skiptypes: Sequence = ()) -> None:
         super(OneHopRSchemaVisitor, self).__init__(skiptypes)
         nodes = set()
         edges = set()
+
         for seschema in rschema.subjects():
             nodes.add((seschema.type, seschema))
+
             for oeschema in rschema.objects(seschema.type):
                 nodes.add((oeschema.type, oeschema))
+
                 if not self.display_rel(rschema, seschema.type, oeschema.type):
                     continue
+
                 edges.add((seschema.type, oeschema.type, rschema))
+
         self._nodes = nodes
         self._edges = edges
 
 
-def schema2dot(schema=None, outputfile=None, skiptypes=(),
-               visitor=None, prophdlr=None, size=None):
+def schema2dot(schema: Optional[_SchemaType] = None, outputfile: Optional[str] = None,
+               skiptypes: Sequence = (), visitor=None, prophdlr=None, size=None) -> Any:
     """write to the output stream a dot graph representing the given schema"""
-    visitor = visitor or FullSchemaVisitor(schema, skiptypes)
+    visitor = visitor or (schema and FullSchemaVisitor(schema, skiptypes))
     prophdlr = prophdlr or SchemaDotPropsHandler(visitor)
+
     if outputfile:
         schemaname = osp.splitext(osp.basename(outputfile))[0]
     else:
         schemaname = 'Schema'
+
     generator = GraphGenerator(DotBackend(schemaname, 'BT',
                                           ratio='compress', size=size,
                                           renderer='dot',
@@ -214,23 +259,31 @@ def schema2dot(schema=None, outputfile=N
                                               # 'polylines':'true',
                                               'sep': '0.2'
                                           }))
+
     return generator.generate(visitor, prophdlr, outputfile)
 
 
-def run():
+def run() -> None:
     """main routine when schema2dot is used as a script"""
     from yams.reader import SchemaLoader
+
     loader = SchemaLoader()
+
     try:
         schema_dir = sys.argv[1]
     except IndexError:
         print("USAGE: schema2dot SCHEMA_DIR [OUTPUT FILE]")
         sys.exit(1)
+
+    outputfile: Optional[str]
+
     if len(sys.argv) > 2:
         outputfile = sys.argv[2]
     else:
         outputfile = None
-    schema = loader.load([schema_dir], 'Test')
+
+    schema = loader.load([schema_dir], 'Test')  # type: ignore # uses old api
+
     schema2dot(schema, outputfile)
 
 
diff --git a/yams/types.py b/yams/types.py
--- a/yams/types.py
+++ b/yams/types.py
@@ -32,6 +32,8 @@ if TYPE_CHECKING:
     _RdefRdefSchemaType = Union[schema.RelationDefinitionSchema, buildobjs.RelationDefinition]
     _RdefType = buildobjs.RelationDefinition
     _SchemaType = schema.Schema
+    _EntitySchemaType = schema.EntitySchema
+    _RelationSchemaType = schema.RelationSchema
     _SubjObjSchema = schema.EntitySchema
     _RtypeType = schema.RelationSchema
 
@@ -39,5 +41,7 @@ else:
     _RdefRdefSchemaType = TypeVar("rdef_rdefschema")
     _RdefType = TypeVar("rdef")
     _SchemaType = TypeVar("schema")
+    _EntitySchemaType = TypeVar("eschema")
+    _RelationSchemaType = TypeVar("rschema")
     _SubjObjSchema = TypeVar("subjobjschema")
     _RtypeType = TypeVar("rtype")



More information about the cubicweb-devel mailing list