Skip to content

Commit d485d86

Browse files
committed
fixes #7
++rebuild method
1 parent 8b147f3 commit d485d86

File tree

2 files changed

+124
-4
lines changed

2 files changed

+124
-4
lines changed

sqlalchemy_mptt/mixins.py

Lines changed: 49 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ def move_after(self, node_id):
8080
session.add(self)
8181

8282
@classmethod
83-
def get_tree(self, session, json=False, json_fields=None):
83+
def get_tree(cls, session, json=False, json_fields=None):
8484
def recursive_node_to_dict(node):
8585
result = {'node': node}
8686
if json:
@@ -93,10 +93,55 @@ def recursive_node_to_dict(node):
9393
result['children'] = children
9494
return result
9595

96-
nodes = session.query(self).filter_by(parent_id=None)\
97-
.order_by(self.tree_id).all()
96+
nodes = session.query(cls).filter_by(parent_id=None)\
97+
.order_by(cls.tree_id).all()
9898
tree = []
99-
for i, node in enumerate(nodes):
99+
for node in nodes:
100100
tree.append(recursive_node_to_dict(node))
101101

102102
return tree
103+
104+
@classmethod
105+
def rebuild_tree(cls, session, tree_id):
106+
session.query(cls).filter_by(tree_id=tree_id)\
107+
.update({cls.left: 0, cls.right: 0, cls.level: 0})
108+
top = session.query(cls).filter_by(parent_id=None)\
109+
.filter_by(tree_id=tree_id).one()
110+
top.left = left = 1
111+
top.right = right = 2
112+
top.level = level = 1
113+
114+
def reqursive(children, left, right, level):
115+
level = level + 1
116+
for i, node in enumerate(children):
117+
same_level_right = children[i-1].right
118+
left = left + 1
119+
120+
if i > 0:
121+
left = left + 1
122+
if same_level_right:
123+
left = same_level_right + 1
124+
125+
right = left + 1
126+
node.left = left
127+
node.right = right
128+
parent = node.parent
129+
130+
j = 0
131+
while parent:
132+
parent.right = right + 1 + j
133+
parent = parent.parent
134+
j += 1
135+
136+
node.level = level
137+
reqursive(node.children, left, right, level)
138+
139+
reqursive(top.children, left, right, level)
140+
141+
@classmethod
142+
def rebuild(cls, session, tree_id=None):
143+
trees = session.query(cls).filter_by(parent_id=None)
144+
if tree_id:
145+
trees = trees.filter_by(tree_id=tree_id)
146+
for tree in trees:
147+
cls.rebuild_tree(session, tree.tree_id)

sqlalchemy_mptt/tests.py

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1050,3 +1050,78 @@ def fields(node):
10501050
tree = Tree.get_tree(self.session, json=True, json_fields=fields)
10511051
self.assertEqual(tree, [{'visible': None, 'children': [{'visible': True, 'children': [{'visible': True, 'id': 3, 'label': '<Node (3)>'}], 'id': 2, 'label': '<Node (2)>'}, {'visible': True, 'children': [{'visible': True, 'id': 5, 'label': '<Node (5)>'}, {'visible': True, 'id': 6, 'label': '<Node (6)>'}], 'id': 4, 'label': '<Node (4)>'}, {'visible': True, 'children': [{'visible': True, 'children': [{'visible': None, 'id': 9, 'label': '<Node (9)>'}], 'id': 8, 'label': '<Node (8)>'}, {'visible': None, 'children': [{'visible': None, 'id': 11, 'label': '<Node (11)>'}], 'id': 10, 'label': '<Node (10)>'}], 'id': 7, 'label': '<Node (7)>'}], 'id': 1, 'label': '<Node (1)>'}, {
10521052
'visible': None, 'children': [{'visible': None, 'children': [{'visible': None, 'id': 14, 'label': '<Node (14)>'}], 'id': 13, 'label': '<Node (13)>'}, {'visible': None, 'children': [{'visible': None, 'id': 16, 'label': '<Node (16)>'}, {'visible': None, 'id': 17, 'label': '<Node (17)>'}], 'id': 15, 'label': '<Node (15)>'}, {'visible': None, 'children': [{'visible': None, 'children': [{'visible': None, 'id': 20, 'label': '<Node (20)>'}], 'id': 19, 'label': '<Node (19)>'}, {'visible': None, 'children': [{'visible': None, 'id': 22, 'label': '<Node (22)>'}], 'id': 21, 'label': '<Node (21)>'}], 'id': 18, 'label': '<Node (18)>'}], 'id': 12, 'label': '<Node (12)>'}])
1053+
1054+
def test_rebuild(self):
1055+
""" level Nested sets tree1
1056+
1 1(1)22
1057+
_______________|___________________
1058+
| | |
1059+
2 2(2)5 6(4)11 12(7)21
1060+
| ^ ^
1061+
3 3(3)4 7(5)8 9(6)10 13(8)16 17(10)20
1062+
| |
1063+
4 14(9)15 18(11)19
1064+
1065+
level Nested sets tree2
1066+
1 1(12)22
1067+
_______________|___________________
1068+
| | |
1069+
2 2(13)5 6(15)11 12(18)21
1070+
| ^ ^
1071+
3 3(14)4 7(16)8 9(17)10 13(19)16 17(21)20
1072+
| |
1073+
4 14(20)15 18(22)19
1074+
1075+
"""
1076+
self.session.query(Tree).update({'lft': 0, 'rgt': 0, 'level': 0})
1077+
Tree.rebuild(self.session, 1)
1078+
self.assertEqual(self.result.all(),
1079+
[(1, 1, 22, 1, None, 1),
1080+
(2, 2, 5, 2, 1, 1),
1081+
(3, 3, 4, 3, 2, 1),
1082+
(4, 6, 11, 2, 1, 1),
1083+
(5, 7, 8, 3, 4, 1),
1084+
(6, 9, 10, 3, 4, 1),
1085+
(7, 12, 21, 2, 1, 1),
1086+
(8, 13, 16, 3, 7, 1),
1087+
(9, 14, 15, 4, 8, 1),
1088+
(10, 17, 20, 3, 7, 1),
1089+
(11, 18, 19, 4, 10, 1),
1090+
1091+
(12, 0, 0, 0, None, 2),
1092+
(13, 0, 0, 0, 12, 2),
1093+
(14, 0, 0, 0, 13, 2),
1094+
(15, 0, 0, 0, 12, 2),
1095+
(16, 0, 0, 0, 15, 2),
1096+
(17, 0, 0, 0, 15, 2),
1097+
(18, 0, 0, 0, 12, 2),
1098+
(19, 0, 0, 0, 18, 2),
1099+
(20, 0, 0, 0, 19, 2),
1100+
(21, 0, 0, 0, 18, 2),
1101+
(22, 0, 0, 0, 21, 2)])
1102+
1103+
Tree.rebuild(self.session)
1104+
self.assertEqual(self.result.all(),
1105+
[(1, 1, 22, 1, None, 1),
1106+
(2, 2, 5, 2, 1, 1),
1107+
(3, 3, 4, 3, 2, 1),
1108+
(4, 6, 11, 2, 1, 1),
1109+
(5, 7, 8, 3, 4, 1),
1110+
(6, 9, 10, 3, 4, 1),
1111+
(7, 12, 21, 2, 1, 1),
1112+
(8, 13, 16, 3, 7, 1),
1113+
(9, 14, 15, 4, 8, 1),
1114+
(10, 17, 20, 3, 7, 1),
1115+
(11, 18, 19, 4, 10, 1),
1116+
1117+
(12, 1, 22, 1, None, 2),
1118+
(13, 2, 5, 2, 12, 2),
1119+
(14, 3, 4, 3, 13, 2),
1120+
(15, 6, 11, 2, 12, 2),
1121+
(16, 7, 8, 3, 15, 2),
1122+
(17, 9, 10, 3, 15, 2),
1123+
(18, 12, 21, 2, 12, 2),
1124+
(19, 13, 16, 3, 18, 2),
1125+
(20, 14, 15, 4, 19, 2),
1126+
(21, 17, 20, 3, 18, 2),
1127+
(22, 18, 19, 4, 21, 2)])

0 commit comments

Comments
 (0)