Skip to content

Commit 081d630

Browse files
Improved naming and functionality for leaves/roots methods
1 parent 6afe2d3 commit 081d630

File tree

1 file changed

+12
-8
lines changed

1 file changed

+12
-8
lines changed

django_postgresql_dag/models.py

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -605,48 +605,52 @@ def ancestors_tree(self):
605605
tree[parent] = parent.ancestors_tree()
606606
return tree
607607

608-
def _get_roots(self, ancestors_tree):
608+
def _roots(self, ancestors_tree):
609609
"""
610610
Works on objects: no queries
611611
"""
612612
if not ancestors_tree:
613613
return set([self])
614614
roots = set()
615615
for ancestor in ancestors_tree:
616-
roots.update(ancestor._get_roots(ancestors_tree[ancestor]))
616+
roots.update(ancestor._roots(ancestors_tree[ancestor]))
617617
return roots
618618

619-
def get_roots(self):
619+
def roots(self):
620620
"""
621621
Returns roots nodes, if any
622622
# ToDo: Modify to use CTE
623623
"""
624624
ancestors_tree = self.ancestors_tree()
625625
roots = set()
626626
for ancestor in ancestors_tree:
627-
roots.update(ancestor._get_roots(ancestors_tree[ancestor]))
627+
roots.update(ancestor._roots(ancestors_tree[ancestor]))
628+
if len(roots) < 1:
629+
roots.add(self)
628630
return roots
629631

630-
def _get_leaves(self, descendants_tree):
632+
def _leaves(self, descendants_tree):
631633
"""
632634
Works on objects: no queries
633635
"""
634636
if not descendants_tree:
635637
return set([self])
636638
leaves = set()
637639
for descendant in descendants_tree:
638-
leaves.update(descendant._get_leaves(descendants_tree[descendant]))
640+
leaves.update(descendant._leaves(descendants_tree[descendant]))
639641
return leaves
640642

641-
def get_leaves(self):
643+
def leaves(self):
642644
"""
643645
Returns leaves nodes, if any
644646
# ToDo: Modify to use CTE
645647
"""
646648
descendants_tree = self.descendants_tree()
647649
leaves = set()
648650
for descendant in descendants_tree:
649-
leaves.update(descendant._get_leaves(descendants_tree[descendant]))
651+
leaves.update(descendant._leaves(descendants_tree[descendant]))
652+
if len(leaves) < 1:
653+
leaves.add(self)
650654
return leaves
651655

652656
@staticmethod

0 commit comments

Comments
 (0)