Fix topological sort logic (#1162)

* Fix topo sort

* Update changelog

Author: CBroz1

CHANGELOG.md

@@ -21,7 +21,8 @@ dj.FreeTable(dj.conn(), "common_session.session_group").drop()
 - Add docstrings to all public methods #1076
 - Update DataJoint to 0.14.2 #1081
 - Allow restriction based on parent keys in `Merge.fetch_nwb()` #1086, #1126
-- Import `datajoint.dependencies.unite_master_parts` -> `topo_sort` #1116, #1137
+- Import `datajoint.dependencies.unite_master_parts` -> `topo_sort` #1116,
+    #1137, #1162
 - Fix bool settings imported from dj config file #1117
 - Allow definition of tasks and new probe entries from config #1074, #1120
 - Enforce match between ingested nwb probe geometry and existing table entry

src/spyglass/utils/dj_graph.py

@@ -17,6 +17,7 @@
 from datajoint.user_tables import TableMeta
 from datajoint.utils import get_master, to_camel_case
 from networkx import (
+    DiGraph,
     NetworkXNoPath,
     NodeNotFound,
     all_simple_paths,
@@ -33,10 +34,36 @@
     unique_dicts,
 )
 
-try:  # Datajoint 0.14.2+ uses topo_sort instead of unite_master_parts
-    from datajoint.dependencies import topo_sort as dj_topo_sort
-except ImportError:
-    from datajoint.dependencies import unite_master_parts as dj_topo_sort
+
+def dj_topo_sort(graph: DiGraph) -> List[str]:
+    """Topologically sort graph.
+
+    Uses datajoint's topo_sort if available, otherwise uses networkx's
+    topological_sort, combined with datajoint's unite_master_parts.
+
+    NOTE: This ordering will impact _hash_upstream, but usage should be
+    consistent before/after a no-transaction populate.
+
+    Parameters
+    ----------
+    graph : nx.DiGraph
+        Directed graph to sort
+
+    Returns
+    -------
+    List[str]
+        List of table names in topological order
+    """
+    __import__("pdb").set_trace()
+    try:  # Datajoint 0.14.2+ uses topo_sort instead of unite_master_parts
+        from datajoint.dependencies import topo_sort
+
+        return topo_sort(graph)
+    except ImportError:
+        from datajoint.dependencies import unite_master_parts
+        from networkx.algorithms.dag import topological_sort
+
+        return unite_master_parts(list(topological_sort(graph)))
 
 
 class Direction(Enum):