"""
Module contains ``order_dicts`` decorator. When used, all dicts declarations are converted to OrderedDict. Example::

    from new_order import order_dicts

    @order_dicts
    def func():
        d1 = {1: 2, 3: 4, 5: 6}
        d2 = {"test": 2 ** 42, "foo": "bar", "baz": d1}
        print("What do we have here? %s" % d1)
        print(d2)


    func()

Output::
    What do we have here? OrderedDict([(1, 2), (3, 4), (5, 6)])
    OrderedDict([('test', 4398046511104), ('foo', 'bar'), ('baz', OrderedDict([(1, 2), (3, 4), (5, 6)]))])

"""
import ast
import inspect
import textwrap


class order_dicts(object):
    """ Decorator to convert all dict literals to OrderedDict.
    """

    def __init__(self, func):
        self.original_function = func
        self.updated_function = None

    def __call__(self, *args, **kwargs):
        if self.updated_function is None:
            source = inspect.getsourcelines(self.original_function)
            source = "".join(source[0])
            source = textwrap.dedent(source)
            code = _convert_dicts_to_ordered(source)
            exec(code, self.original_function.__globals__)

            result = self.original_function.__globals__[self.original_function.__name__]
            while isinstance(result, order_dicts):
                result = result.original_function
            self.updated_function = result
        return self.updated_function(*args, **kwargs)


def _convert_dicts_to_ordered(source):
    root = ast.parse(source)
    import_odict = ast.parse("from collections import OrderedDict")
    root.body.insert(0, import_odict.body[0])
    root = ChangeDictToOrderedTransformer().visit(root)
    root = ast.fix_missing_locations(root)
    return compile(root, "<string>", "exec")


class ChangeDictToOrderedTransformer(ast.NodeTransformer):
    def visit_Dict(self, node):
        self.generic_visit(node)
        tuples = [ast.Tuple(elts=[key, value], ctx=ast.Load())
                  for key, value in zip(node.keys, node.values)]

        od_node = ast.Call(
            func=ast.Name(id="OrderedDict", ctx=ast.Load()),
            args=[ast.Tuple(elts=tuples, ctx=ast.Load())],
            ctx=ast.Load(),
            keywords=[], starargs=None, kwargs=None
        )
        return od_node


# Demo
if __name__ == "__main__":

    @order_dicts
    def func():
        d1 = {1: 2, 3: 4, 5: 6}
        d2 = {"test": 2 ** 42, "foo": "bar", "baz": d1}
        print("What do we have here? %s" % d1)
        print(d2)


    func()