reclosedev · May 13, 2016 13:39 · May 13, 2016 · May 13, 2016
diff --git a/new_order.py b/new_order.py
@@ -0,0 +1,84 @@
+"""
+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()