Merge pull request #231 from dimitri-yatsenko/master

ERD.add_parts

Author: fabiansinz

CHANGELOG.md

@@ -9,3 +9,7 @@
 * 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.
+
+### 0.3.4
+* Added method the `ERD.add_parts` method, which adds the part tables of all tables currently in the ERD.
+* `ERD() + arg` and `ERD() - arg` can now accept relation classes as arg.

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.3"
-__date__ = "July 20, 2016"
+__version__ = "0.3.4"
+__date__ = "July 21, 2016"
 __all__ = ['__author__', '__version__',
            'config', 'conn', 'kill', 'BaseRelation',
            'Connection', 'Heading', 'FreeRelation', 'Not', 'schema',

datajoint/erd.py

@@ -14,7 +14,7 @@
 def _get_tier(table_name):
     try:
         return next(tier for tier in user_relation_classes
-                    if re.fullmatch(tier.tier_regexp, table_name))
+                    if re.fullmatch(tier.tier_regexp, table_name.split('`')[-2]))
     except StopIteration:
         return None
 
@@ -84,6 +84,25 @@ def from_sequence(cls, sequence):
         """
         return functools.reduce(lambda x, y: x+y, map(ERD, sequence))
 
+    def add_parts(self):
+        """
+        Adds to the diagram the part tables of tables already included in the diagram
+        :return:
+        """
+        def is_part(part, master):
+            """
+            :param part:  `database`.`table_name`
+            :param master:   `database`.`table_name`
+            :return: True if part is part of master,
+            """
+            part = [s.strip('`') for s in part.split('.')]
+            master = [s.strip('`') for s in master.split('.')]
+            return master[0] == part[0] and master[1] + '__' == part[1][:len(master[1])+2]
+
+        self = ERD(self)  #  copy
+        self.nodes_to_show.update(n for n in self.nodes() if any(is_part(n, m) for m in self.nodes_to_show))
+        return self
+
     def __add__(self, arg):
         """
         :param arg: either another ERD or a positive integer.
@@ -93,11 +112,14 @@ def __add__(self, arg):
         try:
             self.nodes_to_show.update(arg.nodes_to_show)
         except AttributeError:
-            for i in range(arg):
-                new = nx.algorithms.boundary.node_boundary(self, self.nodes_to_show)
-                if not new:
-                    break
-                self.nodes_to_show.update(new)
+            try:
+                self.nodes_to_show.add(arg.full_table_name)
+            except AttributeError:
+                for i in range(arg):
+                    new = nx.algorithms.boundary.node_boundary(self, self.nodes_to_show)
+                    if not new:
+                        break
+                    self.nodes_to_show.update(new)
         return self
 
     def __sub__(self, arg):
@@ -109,11 +131,14 @@ def __sub__(self, arg):
         try:
             self.nodes_to_show.difference_update(arg.nodes_to_show)
         except AttributeError:
-            for i in range(arg):
-                new = nx.algorithms.boundary.node_boundary(nx.DiGraph(self).reverse(), self.nodes_to_show)
-                if not new:
-                    break
-                self.nodes_to_show.update(new)
+            try:
+                self.nodes_to_show.remove(arg.full_table_name)
+            except AttributeError:
+                for i in range(arg):
+                    new = nx.algorithms.boundary.node_boundary(nx.DiGraph(self).reverse(), self.nodes_to_show)
+                    if not new:
+                        break
+                    self.nodes_to_show.update(new)
         return self
 
     def __mul__(self, arg):
@@ -131,7 +156,7 @@ 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})
+        nx.set_node_attributes(graph, 'node_type', {n: _get_tier(n) for n in graph})
         # relabel nodes to class names
         mapping = {node: lookup_class_name(node, context) for node in graph.nodes()}
         new_names = [mapping.values()]

tests/test_erd.py

@@ -43,7 +43,9 @@ def test_erd_algebra():
         erd1 = erd0 + 3
         erd2 = dj.ERD(E) - 3
         erd3 = erd1 * erd2
+        erd4 = (erd0 + E).add_parts() - B - E
         assert_true(erd0.nodes_to_show == set(cls.full_table_name for cls in [B]))
         assert_true(erd1.nodes_to_show == set(cls.full_table_name for cls in (B, B.C, E, E.F)))
         assert_true(erd2.nodes_to_show == set(cls.full_table_name for cls in (A, B, D, E, L)))
         assert_true(erd3.nodes_to_show == set(cls.full_table_name for cls in (B, E)))
+        assert_true(erd4.nodes_to_show == set(cls.full_table_name for cls in (B.C, E.F)))