improve table.parents and table.children with the option to include foreign key information

Author: dimitri-yatsenko

datajoint/autopopulate.py

@@ -32,25 +32,19 @@ def key_source(self):
                 The default value is the join of the parent relations.
                 Users may override to change the granularity or the scope of populate() calls.
         """
-        def parent_gen(self):
-            if self.target.full_table_name not in self.connection.dependencies:
-                self.connection.dependencies.load()
-            for parent_name, fk_props in self.target.parents(primary=True).items():
-                if not parent_name.isdigit():  # simple foreign key
-                    yield FreeTable(self.connection, parent_name).proj()
-                else:
-                    grandparent = list(self.connection.dependencies.in_edges(parent_name))[0][0]
-                    yield FreeTable(self.connection, grandparent).proj(**{
-                        attr: ref for attr, ref in fk_props['attr_map'].items() if ref != attr})
+        def _rename_attributes(table, props):
+            return (table.proj(
+                **{attr: ref for attr, ref in props['attr_map'].items() if attr != ref})
+                if props['aliased'] else table)
 
         if self._key_source is None:
-            parents = parent_gen(self)
-            try:
-                self._key_source = next(parents)
-            except StopIteration:
-                raise DataJointError('A relation must have primary dependencies for auto-populate to work') from None
-            for q in parents:
-                self._key_source *= q
+            parents = self.target.parents(primary=True, as_objects=True, foreign_key_info=True)
+            if not parents:
+                raise DataJointError(
+                    'A relation must have primary dependencies for auto-populate to work') from None
+            self._key_source = _rename_attributes(*parents[0])
+            for q in parents[1:]:
+                self._key_source *= _rename_attributes(*q)
         return self._key_source
 
     def make(self, key):

datajoint/table.py

@@ -11,7 +11,7 @@
 from .declare import declare, alter
 from .expression import QueryExpression
 from . import blob
-from .utils import user_choice
+from .utils import user_choice, OrderedDict
 from .heading import Heading
 from .errors import DuplicateError, AccessError, DataJointError, UnknownAttributeError
 from .version import __version__ as version
@@ -40,8 +40,7 @@ class Table(QueryExpression):
     @property
     def heading(self):
         """
-        Returns the table heading. If the table is not declared, attempts to declare it and return heading.
-        :return: table heading
+        :return: table heading. If the table is not declared, attempts to declare it first.
         """
         if self._heading is None:
             self._heading = Heading()  # instance-level heading
@@ -53,7 +52,8 @@ def heading(self):
     def declare(self, context=None):
         """
         Declare the table in the schema based on self.definition.
-        :param context: the context for foreign key resolution. If None, foreign keys are not allowed.
+        :param context: the context for foreign key resolution. If None, foreign keys are
+            not allowed.
         """
         if self.connection.in_transaction:
             raise DataJointError('Cannot declare new tables inside a transaction, '
@@ -116,38 +116,59 @@ def get_select_fields(self, select_fields=None):
         """
         return '*' if select_fields is None else self.heading.project(select_fields).as_sql
 
-    def parents(self, primary=None, as_objects=False):
+    def parents(self, primary=None, as_objects=False, foreign_key_info=False):
         """
         :param primary: if None, then all parents are returned. If True, then only foreign keys composed of
-            primary key attributes are considered.  If False, the only foreign keys including at least one non-primary
-            attribute are considered.
-        :param as_objects: if False (default), the output is a dict describing the foreign keys. If True, return table objects.
-        :return: dict of tables referenced with self's foreign keys  or list of table objects if as_objects=True
-        """
-        parents = self.connection.dependencies.parents(self.full_table_name, primary)
+            primary key attributes are considered.  If False, return foreign keys including at least one
+            secondary attribute.
+        :param as_objects: if False, return table names. If True, return table objects.
+        :param foreign_key_info: if True, each element in result also includes foreign key info.
+        :return: list of parents as table names or table objects
+            with (optional) foreign key information.
+        """
+        get_edge = self.connection.dependencies.parents
+        nodes = [next(iter(get_edge(name).items())) if name.isdigit() else (name, props)
+                 for name, props in get_edge(self.full_table_name, primary).items()]
         if as_objects:
-            parents = [FreeTable(self.connection, c) for c in parents]
-        return parents
+            nodes = [(FreeTable(self.connection, name), props) for name, props in nodes]
+        if not foreign_key_info:
+            nodes = [name for name, props in nodes]
+        return nodes
 
-    def children(self, primary=None, as_objects=False):
+    def children(self, primary=None, as_objects=False, foreign_key_info=False):
         """
         :param primary: if None, then all children are returned. If True, then only foreign keys composed of
-            primary key attributes are considered.  If False, the only foreign keys including at least one non-primary
-            attribute are considered.
-        :param as_objects: if False (default), the output is a dict describing the foreign keys. If True, return table objects.
-        :return: dict of tables with foreign keys referencing self or list of table objects if as_objects=True
-        """
-        nodes = dict((next(iter(self.connection.dependencies.children(k).items())) if k.isdigit() else (k, v))
-                     for k, v in self.connection.dependencies.children(self.full_table_name, primary).items())
+            primary key attributes are considered.  If False, return foreign keys including at least one
+            secondary attribute.
+        :param as_objects: if False, return table names. If True, return table objects.
+        :param foreign_key_info: if True, each element in result also includes foreign key info.
+        :return: list of children as table names or table objects
+            with (optional) foreign key information.
+        """
+        get_edge = self.connection.dependencies.children
+        nodes = [next(iter(get_edge(name).items())) if name.isdigit() else (name, props)
+                 for name, props in get_edge(self.full_table_name, primary).items()]
         if as_objects:
-            nodes = [FreeTable(self.connection, c) for c in nodes]
+            nodes = [(FreeTable(self.connection, name), props) for name, props in nodes]
+        if not foreign_key_info:
+            nodes = [name for name, props in nodes]
         return nodes
 
     def descendants(self, as_objects=False):
-        nodes = [node for node in self.connection.dependencies.descendants(self.full_table_name) if not node.isdigit()]
-        if as_objects:
-            nodes = [FreeTable(self.connection, c) for c in nodes]
-        return nodes
+        """
+        :param as_objects: False - a list of table names; True - a list of table objects.
+        :return: list of tables descendants in topological order.
+        """
+        return [FreeTable(self.connection, node) if as_objects else node
+                for node in self.connection.dependencies.descendants(self.full_table_name) if not node.isdigit()]
+
+    def ancestors(self, as_objects=False):
+        """
+        :param as_objects: False - a list of table names; True - a list of table objects.
+        :return: list of tables ancestors in topological order.
+        """
+        return [FreeTable(self.connection, node) if as_objects else node
+                for node in self.connection.dependencies.ancestors(self.full_table_name) if not node.isdigit()]
 
     def parts(self, as_objects=False):
         """
@@ -156,13 +177,7 @@ def parts(self, as_objects=False):
         """
         nodes = [node for node in self.connection.dependencies.nodes
                  if not node.isdigit() and node.startswith(self.full_table_name[:-1] + '__')]
-        if as_objects:
-            nodes = [FreeTable(self.connection, c) for c in nodes]
-        return nodes
-
-    def ancestors(self, as_objects=False):
-        return [FreeTable(self.connection, node) if as_objects else node
-                for node in self.connection.dependencies.ancestors(self.full_table_name) if not node.isdigit()]
+        return [FreeTable(self.connection, c) for c in nodes] if as_objects else nodes
 
     @property
     def is_declared(self):
@@ -525,7 +540,7 @@ def describe(self, context=None, printout=True):
             del frame
         if self.full_table_name not in self.connection.dependencies:
             self.connection.dependencies.load()
-        parents = self.parents()
+        parents = self.parents(foreign_key_info=True)
         in_key = True
         definition = ('# ' + self.heading.table_info['comment'] + '\n'
                       if self.heading.table_info['comment'] else '')
@@ -538,11 +553,10 @@ def describe(self, context=None, printout=True):
                 in_key = False
             attributes_thus_far.add(attr.name)
             do_include = True
-            for parent_name, fk_props in list(parents.items()):  # need list() to force a copy
+            for parent_name, fk_props in parents:
                 if attr.name in fk_props['attr_map']:
                     do_include = False
                     if attributes_thus_far.issuperset(fk_props['attr_map']):
-                        parents.pop(parent_name)
                         # foreign key properties
                         try:
                             index_props = indexes.pop(tuple(fk_props['attr_map']))
@@ -552,19 +566,19 @@ def describe(self, context=None, printout=True):
                             index_props = [k for k, v in index_props.items() if v]
                             index_props = ' [{}]'.format(', '.join(index_props)) if index_props else ''
 
-                        if not parent_name.isdigit():
+                        if not fk_props['aliased']:
                             # simple foreign key
                             definition += '->{props} {class_name}\n'.format(
                                 props=index_props,
                                 class_name=lookup_class_name(parent_name, context) or parent_name)
                         else:
                             # projected foreign key
-                            parent_name = list(self.connection.dependencies.in_edges(parent_name))[0][0]
-                            lst = [(attr, ref) for attr, ref in fk_props['attr_map'].items() if ref != attr]
                             definition += '->{props} {class_name}.proj({proj_list})\n'.format(
                                 props=index_props,
                                 class_name=lookup_class_name(parent_name, context) or parent_name,
-                                proj_list=','.join('{}="{}"'.format(a, b) for a, b in lst))
+                                proj_list=','.join(
+                                    '{}="{}"'.format(attr, ref)
+                                    for attr, ref in fk_props['attr_map'].items() if ref != attr))
                             attributes_declared.update(fk_props['attr_map'])
             if do_include:
                 attributes_declared.add(attr.name)

tests/test_fetch.py

@@ -201,7 +201,7 @@ def test_fetch1_step3(self):
         self.lang.fetch1('name')
 
     def test_decimal(self):
-        """Tests that decimal fields are correctly fetched and used in restrictions, see issue #334"""
+        """ Tests that decimal fields are correctly fetched and used in restrictions, see issue #334"""
         rel = schema.DecimalPrimaryKey()
         rel.insert1([decimal.Decimal('3.1415926')])
         keys = rel.fetch()

tests/test_foreign_keys.py

@@ -1,4 +1,4 @@
-from nose.tools import assert_equal, assert_false, assert_true, raises
+from nose.tools import assert_equal, assert_false, assert_true
 from datajoint.declare import declare
 
 from . import schema_advanced