from theano import gof
from theano import compile
from theano.gof import TopoOptimizer
from theano.typed_list.basic import (Reverse,
                    Append, Extend, Insert, Remove)


@gof.local_optimizer([Append, Extend, Insert, Reverse, Remove], inplace=True)
def typed_list_inplace_opt(node):
    if isinstance(node.op, (Append, Extend, Insert, Reverse, Remove)) \
                                        and not node.op.inplace:

        new_op = node.op.__class__(
                    inplace=True)
        new_node = new_op(*node.inputs)
        return [new_node]
    return False
compile.optdb.register('typed_list_inplace_opt',
                       TopoOptimizer(typed_list_inplace_opt,
    failure_callback=TopoOptimizer.warn_inplace), 60,
                       'fast_run', 'inplace')
