Skip to content

Commit 524b497

Browse files
Improved edge-related methods on Node
1 parent 007627e commit 524b497

File tree

1 file changed

+23
-21
lines changed

1 file changed

+23
-21
lines changed

django_postgresql_dag/models.py

Lines changed: 23 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -226,56 +226,58 @@ def clan(self):
226226
"""
227227
return self.filter_order_ids(self.clan_ids())
228228

229-
def descendants_edges_set(self, cached_results=None):
229+
def descendants_edges(self, cached_results=None):
230230
"""
231-
Returns a set of descendants edges
231+
Returns a queryset of descendants edges
232232
# ToDo: Modify to use CTE
233233
"""
234234
if cached_results is None:
235235
cached_results = dict()
236236
if self.id in cached_results.keys():
237237
return cached_results[self.id]
238238
else:
239-
res = set()
239+
edge_set = set()
240240
for f in self.children.all():
241-
res.add(edge_model.objects.get(parent=self.id, child=f.id))
242-
res.update(f.descendants_edges_set(cached_results=cached_results))
243-
cached_results[self.id] = res
244-
return res
241+
edge_set.add(edge_model.objects.get(parent=self.id, child=f.id))
242+
edge_set.update(f.descendants_edges(cached_results=cached_results))
243+
cached_results[self.id] = edge_set
244+
out = [x.id for x in edge_set]
245+
return _filter_order(edge_model.objects, "pk", out)
245246

246-
def ancestors_edges_set(self, cached_results=None):
247+
def ancestors_edges(self, cached_results=None):
247248
"""
248-
Returns a set of ancestors edges
249+
Returns a queryset of ancestors edges
249250
# ToDo: Modify to use CTE
250251
"""
251252
if cached_results is None:
252253
cached_results = dict()
253254
if self in cached_results.keys():
254255
return cached_results[self.id]
255256
else:
256-
res = set()
257+
edge_set = set()
257258
for f in self.parents.all():
258-
res.add(edge_model.objects.get(child=self.id, parent=f.id))
259-
res.update(f.ancestors_edges_set(cached_results=cached_results))
260-
cached_results[self.id] = res
261-
return res
259+
edge_set.add(edge_model.objects.get(child=self.id, parent=f.id))
260+
edge_set.update(f.ancestors_edges(cached_results=cached_results))
261+
cached_results[self.id] = edge_set
262+
out = [x.id for x in edge_set]
263+
return _filter_order(edge_model.objects, "pk", out)
262264

263-
def clan_edges_set(self):
265+
def clan_edges(self):
264266
"""
265-
Returns a set of all edges associated with a given node
267+
Returns a queryset of all edges associated with a given node
266268
"""
267269
edges = set()
268-
edges.update(self.descendants_edges_set())
269-
edges.update(self.ancestors_edges_set())
270+
edges.update(self.descendants_edges())
271+
edges.update(self.ancestors_edges())
270272
return edges
271273

272274
def path_ids_list(
273275
self, target_node, directional=True, max_depth=20, max_paths=1
274276
):
275277
"""
276-
Returns a list of paths from self to target node, in either direction.
277-
The resulting lists are always sorted from top of graph, downward,
278-
regardless of the relative position of starting and ending nodes.
278+
Returns a list of paths from self to target node, optionally in either
279+
direction. The resulting lists are always sorted from root-side, toward
280+
leaf-side, regardless of the relative position of starting and ending nodes.
279281
280282
By default, returns only one shortest path, but additional paths
281283
can be included by setting the max_paths argument.

0 commit comments

Comments
 (0)