Add sum_tree helper

trytond/CHANGELOG

@@ -1,3 +1,4 @@
+* Add sum_tree helper
 * Convert sum attribute into boolean
 * Store records selection on export
 

trytond/doc/ref/models.rst

@@ -1100,6 +1100,14 @@ name by concatenating each parent names using it as separator and
    Helper method that checks if there is no recursion in the tree defined by
    :meth:`tree`.
 
+
+.. function:: sum_tree(records, values[, parent])
+
+   Helper method to sum ``values`` following up ``records`` tree using
+   ``parent``.
+   ``records`` must contain a complete branch of the tree.
+   ``values`` is a dictionary with record id as key.
+
 avatar_mixin
 ============
 

trytond/trytond/model/__init__.py

@@ -14,7 +14,7 @@ from .modelview import ModelView
 from .multivalue import MultiValueMixin, ValueMixin
 from .order import sequence_ordered, sort
 from .symbol import SymbolMixin
-from .tree import tree
+from .tree import sum_tree, tree
 from .union import UnionMixin
 from .workflow import Workflow
 
@@ -23,4 +23,4 @@ __all__ = ['Model', 'ModelView', 'ModelStorage', 'ModelSingleton', 'ModelSQL',
     'Workflow', 'DictSchemaMixin', 'MatchMixin', 'UnionMixin', 'dualmethod',
     'MultiValueMixin', 'ValueMixin', 'SymbolMixin', 'DigitsMixin',
     'EvalEnvironment', 'sequence_ordered', 'sort', 'DeactivableMixin', 'tree',
-    'avatar_mixin']
+    'sum_tree', 'avatar_mixin']

trytond/trytond/model/tree.py

@@ -146,3 +146,27 @@ def tree(parent='parent', name='name', separator=None):
                 visited.update(walked)
 
     return TreeMixin
+
+
+def sum_tree(records, values, parent='parent'):
+    "Sum up values following records tree"
+    result = values.copy()
+    parents = {
+        int(r): int(getattr(r, parent)) for r in records if getattr(r, parent)}
+    records = set(map(int, records))
+    leafs = records - set(parents.values())
+    while leafs:
+        for leaf in leafs:
+            records.remove(leaf)
+            parent = parents.get(leaf)
+            if parent:
+                result[parent] += result[leaf]
+        next_leafs = set(records)
+        for record in records:
+            parent = parents.get(record)
+            if not parent:
+                continue
+            if parent in next_leafs and parent in records:
+                next_leafs.remove(parent)
+        leafs = next_leafs
+    return result

trytond/trytond/tests/test_tree.py

@@ -2,7 +2,9 @@
 # this repository contains the full copyright notices and license terms.
 
 import unittest
+from collections import defaultdict
 
+from trytond.model import sum_tree
 from trytond.model.exceptions import DomainValidationError, RecursionError
 from trytond.pool import Pool
 from trytond.tests.test_tryton import activate_module, with_transaction
@@ -429,3 +431,28 @@ class TreeMixinTestCase(unittest.TestCase):
         with self.assertRaises(RecursionError):
             parent1.parents = [child]
             parent1.save()
+
+
+class TreeTestCase(unittest.TestCase):
+    "Test Tree"
+
+    @with_transaction()
+    def test_sum_tree(self):
+        pool = Pool()
+        Model = pool.get('test.tree')
+        records = [
+            Model(1, parent=None),
+            Model(2, parent=1),
+            Model(3, parent=1),
+            Model(4, parent=2),
+            Model(5, parent=2),
+            ]
+        values = defaultdict(lambda: 1)
+        self.assertEqual(
+            dict(sum_tree(records, values)), {
+                5: 1,
+                4: 1,
+                3: 1,
+                2: 3,
+                1: 5,
+                })