Merge pull request #230 from dimitri-yatsenko/master

Improved ERD

Author: fabiansinz

CHANGELOG.md

@@ -3,3 +3,9 @@
 * Fixed issue #223:  `insert` can insert relations without fetching
 * ERD() now takes the `context` argument, which specifies in which context to look for classes. The default is taken from the argument (schema or relation).
 * ERD.draw() no longer has the `prefix` argument: class names are shown as found in the context.
+
+### 0.3.3
+* Suppressed warnings (redirected them to logging).  Previoiusly, scipy would throw warnings in ERD, for example.
+* Added ERD.from_sequence as a shortcut to combining the ERDs of multiple sources
+* ERD() no longer text the context argument.
+* ERD.draw() now takes an optional context argument.  By default uses the caller's locals.

datajoint/__init__.py

@@ -16,8 +16,8 @@
 import os
 
 __author__ = "Dimitri Yatsenko, Edgar Walker, and Fabian Sinz at Baylor College of Medicine"
-__version__ = "0.3.2"
-__date__ = "July 18, 2016"
+__version__ = "0.3.3"
+__date__ = "July 20, 2016"
 __all__ = ['__author__', '__version__',
            'config', 'conn', 'kill', 'BaseRelation',
            'Connection', 'Heading', 'FreeRelation', 'Not', 'schema',
@@ -27,6 +27,8 @@
 
 print('DataJoint', __version__, '('+__date__+')')
 
+logging.captureWarnings(True)
+
 
 class key:
     """

datajoint/base_relation.py

@@ -355,7 +355,7 @@ def show_definition(self):
                     do_include = False
                 if attributes_thus_far.issuperset(primary_key):
                     parents.pop(parent)
-                    definition += '-> ' + lookup_class_name(self.context, parent) + '\n'
+                    definition += '-> ' + lookup_class_name(parent, self.context) + '\n'
             if do_include:
                 definition += '%-20s : %-28s # %s\n' % (
                     attr.name if attr.default is None else '%s = %s' % (attr.name, attr.default),
@@ -364,33 +364,39 @@ def show_definition(self):
         return definition
 
 
-def lookup_class_name(context, name, levels=4):
+def lookup_class_name(name, context, depth=3):
     """
     given the name of another table in the form `database`.`table_name`, find its class in the context
-    :param context: dictionary representing the namespace
     :param name: `database`.`table_name`
-    :param levels: the limit on how deep to go into imported modules, helps avoid infinite recursion.
-    :return: class name found in the context.
+    :param context: dictionary representing the namespace
+    :param depth: search depth into imported modules, helps avoid infinite recursion.
+    :return: class name found in the context or original name if not found
     """
-    for member_name, member in context.items():
-        if inspect.isclass(member) and issubclass(member, BaseRelation):
-            if member.full_table_name == name:
-                return member_name
-            # check part tables
-            try:
-                parts = member._ordered_class_members
-            except AttributeError:
-                pass  # not a UserRelation
-            else:
-                for part in parts:
-                    part = getattr(member, part)
-                    if inspect.isclass(part) and issubclass(part, BaseRelation) and part.full_table_name == name:
-                        return member_name + '.' + part.__name__
-        elif levels and inspect.ismodule(member) and member.__name__ != 'datajoint':
-            candidate_name = lookup_class_name(dict(inspect.getmembers(member)), name, levels-1)
-            if candidate_name != name:
-                return member_name + '.' + candidate_name
-    return name
+    # breadth-first search
+    nodes = [dict(context=context, context_name='', depth=depth)]
+    while nodes:
+        node = nodes.pop(0)
+        for member_name, member in node['context'].items():
+            if inspect.isclass(member) and issubclass(member, BaseRelation):
+                if member.full_table_name == name:   # found it!
+                    return '.'.join([node['context_name'],  member_name]).lstrip('.')
+                try:  # look for part tables
+                    parts = member._ordered_class_members
+                except AttributeError:
+                    pass  # not a UserRelation -- cannot have part tables.
+                else:
+                    for part in (getattr(member, p) for p in parts):
+                        if inspect.isclass(part) and issubclass(part, BaseRelation) and part.full_table_name == name:
+                            return '.'.join([node['context_name'], member_name, part.__name__]).lstrip('.')
+            elif node['depth'] > 0 and inspect.ismodule(member) and member.__name__ != 'datajoint':
+                try:
+                    nodes.append(
+                        dict(context=dict(inspect.getmembers(member)),
+                             context_name=node['context_name'] + '.' + member_name,
+                             depth=node['depth']-1))
+                except ImportError:
+                    pass  # could not import, so do not attempt
+    return name  # return original name if class is not found
 
 
 class FreeRelation(BaseRelation):

datajoint/connection.py

@@ -62,6 +62,7 @@ def __init__(self, host, user, passwd, init_fun=None):
             port = config['database.port']
         self.conn_info = dict(host=host, port=port, user=user, passwd=passwd)
         self.init_fun = init_fun
+        print("Connecting {user}@{host}:{port}".format(**self.conn_info))
         self.connect()
         if self.is_connected:
             logger.info("Connected {user}@{host}:{port}".format(**self.conn_info))

datajoint/erd.py

@@ -1,12 +1,13 @@
 import networkx as nx
 import re
 import numpy as np
+import functools
 from scipy.optimize import basinhopping
 import itertools
-import inspect
 from . import Manual, Imported, Computed, Lookup, Part, DataJointError
 from .base_relation import lookup_class_name
 
+
 user_relation_classes = (Manual, Lookup, Computed, Imported, Part)
 
 
@@ -23,10 +24,8 @@ class ERD(nx.DiGraph):
     Entity relationship diagram.
 
     Usage:
-    >>>  erd = Erd(source, context=None)
+    >>>  erd = Erd(source)
     source can be a base relation object, a base relation class, a schema, or a module that has a schema.
-    context is the namespace in which to search for classes. If None (default), looks
-    in source.context or source.schema.context
 
     >>> erd.draw()
     draws the diagram using pyplot
@@ -39,12 +38,11 @@ class ERD(nx.DiGraph):
     Note that erd + 1 - 1  may differ from erd - 1 + 1 and so forth.
     Only those tables that are loaded in the connection object are displayed
     """
-    def __init__(self, source, context=None):
+    def __init__(self, source):
 
         if isinstance(source, ERD):
             # copy constructor
             self.nodes_to_show = set(source.nodes_to_show)
-            self._context = source._context
             super().__init__(source)
             return
 
@@ -57,18 +55,6 @@ def __init__(self, source, context=None):
             except AttributeError:
                 raise DataJointError('Could not find database connection in %s' % repr(source[0]))
 
-        if context is None:
-            # find context in the source
-            try:
-                context = source.context
-            except AttributeError:
-                try:
-                    context = source.schema.context
-                except AttributeError:
-                    raise DataJointError('Could not find context in %s' % repr(source[0]))
-
-        self._context = context
-
         # initialize graph from dependencies
         connection.dependencies.load()
         super().__init__(connection.dependencies)
@@ -89,6 +75,15 @@ def __init__(self, source, context=None):
                 if node.startswith('`%s`' % database):
                     self.nodes_to_show.add(node)
 
+    @classmethod
+    def from_sequence(cls, sequence):
+        """
+        The join ERD for all objects in sequence
+        :param sequence: a sequence (e.g. list, tuple)
+        :return: ERD(arg1) + ... + ERD(argn)
+        """
+        return functools.reduce(lambda x, y: x+y, map(ERD, sequence))
+
     def __add__(self, arg):
         """
         :param arg: either another ERD or a positive integer.
@@ -131,25 +126,42 @@ def __mul__(self, arg):
         self.nodes_to_show.intersection_update(arg.nodes_to_show)
         return self
 
-    def _make_graph(self):
+    def _make_graph(self, context):
         """
         Make the self.graph - a graph object ready for drawing
         """
         graph = nx.DiGraph(self).subgraph(self.nodes_to_show)
         nx.set_node_attributes(graph, 'node_type', {n: _get_tier(n.split('`')[-2]) for n in graph})
         # relabel nodes to class names
-        mapping = {node: lookup_class_name(self._context, node) for node in graph.nodes()}
+        mapping = {node: lookup_class_name(node, context) for node in graph.nodes()}
         new_names = [mapping.values()]
         if len(new_names) > len(set(new_names)):
             raise DataJointError('Some classes have identical names. The ERD cannot be plotted.')
         nx.relabel_nodes(graph, mapping, copy=False)
         return graph
 
-    def draw(self, pos=None, layout=None, font_scale=1.2, **layout_options):
+    def draw(self, pos=None, layout=None, context=None, font_scale=1.2, **layout_options):
+        """
+        Draws the graph of dependencies.
+        :param pos: dict with positions for every node.  If None, then layout is called.
+        :param layout: the graph layout function. If None, then self._layout is used.
+        :param context: the context in which to look for the class names.  If None, the caller's context is used.
+        :param font_scale: the scalar used to scale all the fonts.
+        :param layout_options:  kwargs passed into the layout function.
+        """
         if not self.nodes_to_show:
             print('There is nothing to plot')
             return
-        graph = self._make_graph()
+        if context is None:
+            # get the caller's locals()
+            import inspect
+            frame = inspect.currentframe()
+            try:
+                context = frame.f_back.f_locals
+            finally:
+                del frame
+
+        graph = self._make_graph(context)
         if pos is None:
             pos = (layout if layout else self._layout)(graph, **layout_options)
         import matplotlib.pyplot as plt
@@ -158,7 +170,7 @@ def draw(self, pos=None, layout=None, font_scale=1.2, **layout_options):
         edge_styles = ['solid' if e[2]['primary'] else 'dashed' for e in edge_list]
         nx.draw_networkx_edges(graph, pos=pos, edgelist=edge_list, style=edge_styles, alpha=0.2)
 
-        label_props = { # http://matplotlib.org/examples/color/named_colors.html
+        label_props = {  # http://matplotlib.org/examples/color/named_colors.html
             None: dict(bbox=dict(boxstyle='round,pad=0.1', facecolor='yellow', alpha=0.3), size=round(font_scale*8)),
             Manual: dict(bbox=dict(boxstyle='round,pad=0.1', edgecolor='white', facecolor='darkgreen', alpha=0.3), size=round(font_scale*10)),
             Lookup: dict(bbox=dict(boxstyle='round,pad=0.1', edgecolor='white', facecolor='gray', alpha=0.2), size=round(font_scale*8)),

tests/test_erd.py

@@ -32,7 +32,7 @@ def test_dependencies():
     @staticmethod
     def test_erd():
         erd = dj.ERD(schema)
-        graph = erd._make_graph()
+        graph = erd._make_graph(schema.context)
         assert_true(set(cls.__name__ for cls in (A, B, D, E, L)).issubset(graph.nodes()))
         pos = erd._layout(graph)
         assert_true(set(cls.__name__ for cls in (A, B, D, E, L)).issubset(pos.keys()))