New experimental way to handle trees in deap

gp2.py

import gp
import random

__type__ = None

class PrimitiveTree(list):
    def __init__(self, content):
        self[:] = content

    @property
    def size(self):
        return len(self)

    @property
    def height(self):
        stack = [0]
        max_depth = 0
        for elem in self:
            depth = stack.pop()
            max_depth = max(max_depth, depth)
            stack.extend([depth+1] * elem.arity)
        return max_depth

    @property
    def root(self):
        return self[0]

    @property
    def iter_leaf(self):
        for elem in self:
            if elem.arity == 0:
                yield elem

    @property
    def iter_leaf_idx(self):
        for idx, elem in enumerate(self):
            if elem.arity == 0:
                yield idx

    def searchSubtree(self, begin):
        end = begin + 1        
        total = self[begin].arity
        while total > 0:
            total += self[end].arity - 1
            end += 1
        return slice(begin, end)

    def setSubtree(self, begin, subtree):
        slice_ = self.searchSubtree(begin)
        self[slice_] = subtree

def condition(height, depth):
    """Expression generation stops when the depth is equal to height."""
    return depth == height

def generate(pset, min_, max_, condition=condition, type_=__type__):
    """Generate a Tree as a list of list. The tree is build
    from the root to the leaves, and it stop growing when the
    condition is fulfilled. 

    :param pset: A primitive set from wich to select primitives of the trees.
    :param min_: Minimum height of the produced trees.
    :param max_: Maximum Height of the produced trees.
    :param condition: The condition is a function that kes two arguments, 
                      the height of the tree to build and the current 
                      depth in the tree.
    :param type_: The type that should return the tree when called, when
                  :obj:`None` (default) no return type is enforced.
    :returns: A grown tree with leaves at possibly different depths 
              dependending on the condition function.
    """
    expr = []
    height = random.randint(min_, max_)
    stack = [(0, type_)]
    while len(stack) != 0:
        depth, type_ = stack.pop()
        if condition(height, depth):
            term = random.choice(pset.terminals[type_])
            expr.append(term())
        else:
            prim = random.choice(pset.primitives[type_])
            expr.append(prim)
            for arg in prim.args:
                stack.append((depth+1, arg))
    return expr