Feat(Tree): Add organise for AssocialTree

mapytex/calculus/core/tree.py

@@ -1099,28 +1099,43 @@ class AssocialTree(Tree):
 
         :example:
 
-        >>> t = Tree.from_str("3*4+2", convert_to_mo=False)
+        >>> t = AssocialTree.from_str("3*4+2", convert_to_mo=False)
         >>> print(t)
         +
          > *
          | > 3
          | > 4
          > 2
-        >>> print(t.map_on_leaf(lambda x:2*x))
+        >>> print(t.map_on_leaf(lambda x:10*x))
+        Traceback (most recent call last):
+        ...
+        TypeError: unsupported operand type(s) for *: 'int' and 'AssocialTree'
+        >>> def try_multiply_ten(x):
+        ...    try:
+        ...        return x*10
+        ...    except:
+        ...        return x
+        >>> print(t.map_on_leaf(try_multiply_ten))
         +
          > *
-         | > 6
+         | > 3
-         | > 8
+         | > 4
-         > 4
+         > 20
 
         """
         try:
-            left_applied = self.left_value.map_on_leaf(function)
+            if self.left_value.node == self.node:
+                left_applied = self.left_value.map_on_leaf(function)
+            else:
+                left_applied = function(self.left_value)
         except AttributeError:
             left_applied = function(self.left_value)
 
         try:
-            right_applied = self.right_value.map_on_leaf(function)
+            if self.right_value.node == self.node:
+                right_applied = self.right_value.map_on_leaf(function)
+            else:
+                right_applied = function(self.right_value)
         except AttributeError:
             right_applied = function(self.right_value)
 
@@ -1296,6 +1311,145 @@ class AssocialTree(Tree):
         t = Tree.from_list(self.node, balanced_leafs)
         return t
 
+    def organise_by(self, 
+                    signature=lambda x: type(x),
+                    recursive=True,
+                    exclude_nodes=[]):
+        """ Reoganise AssocialTree base on self order and groups by signature
+
+        :param signature: grouping function (default type)
+        :param recursive: treat nested AssocialTree the same way (default True)
+        :param exclude_nodes: do not organise trees with thoses nodes (default [])
+
+        :return: an organise version of self
+
+        :example:
+        >>> t = AssocialTree.from_list('+', [3, 4.1, 'y', 55, 2.3, 'x'])
+        >>> print(t)
+        +
+         > +
+         | > 3
+         | > +
+         | | > 4.1
+         | | > y
+         > +
+         | > 55
+         | > +
+         | | > 2.3
+         | | > x
+        >>> print(t.organise_by())
+        +
+         > +
+         | > 3
+         | > 55
+         > +
+         | > +
+         | | > 4.1
+         | | > 2.3
+         | > +
+         | | > y
+         | | > x
+        >>> t = AssocialTree.from_list('+', [1, 'x', 3, 'y'])
+        >>> T = AssocialTree.from_list('*', [5, 'v', 6, 'w',  t])
+        >>> print(T)
+        *
+         > *
+         | > 5
+         | > v
+         > *
+         | > 6
+         | > *
+         | | > w
+         | | > +
+         | | | > +
+         | | | | > 1
+         | | | | > x
+         | | | > +
+         | | | | > 3
+         | | | | > y
+        >>> print(T.organise_by())
+        *
+         > *
+         | > 5
+         | > 6
+         > *
+         | > *
+         | | > v
+         | | > w
+         | > +
+         | | > +
+         | | | > 1
+         | | | > 3
+         | | > +
+         | | | > x
+         | | | > y
+        >>> print(T.organise_by(recursive=False))
+        *
+         > *
+         | > 5
+         | > 6
+         > *
+         | > *
+         | | > v
+         | | > w
+         | > +
+         | | > +
+         | | | > 1
+         | | | > x
+         | | > +
+         | | | > 3
+         | | | > y
+        >>> print(T.organise_by(exclude_nodes=['*']))
+        *
+         > *
+         | > 5
+         | > v
+         > *
+         | > 6
+         | > *
+         | | > w
+         | | > +
+         | | | > +
+         | | | | > 1
+         | | | | > 3
+         | | | > +
+         | | | | > x
+         | | | | > y
+        """
+        if self.node not in exclude_nodes:
+            groups = {}
+            for leaf in self.get_leafs():
+                if signature(leaf) in groups:
+                    groups[signature(leaf)].append(leaf)
+                else:
+                    groups[signature(leaf)] = [leaf]
+
+            subtrees = []
+            for group in groups.values():
+                try:
+                    subtrees.append(Tree.from_list(self.node, group))
+                except ValueError:
+                    subtrees.append(*group)
+            if len(subtrees) > 1:
+                tree = Tree.from_list(self.node, subtrees)
+            else:
+                tree = subtrees[0]
+        else:
+            tree = self
+
+        if recursive:
+
+            def contaminate_organise(leaf):
+                try:
+                    return leaf.organise_by(signature, recursive, exclude_nodes)
+                except AttributeError:
+                    return leaf
+
+            return AssocialTree.from_any_tree(tree).\
+                map_on_leaf(contaminate_organise)
+
+        return tree
+
 # -----------------------------
 # Reglages pour 'vim'
 # vim:set autoindent expandtab tabstop=4 shiftwidth=4: