Modify AST line order

gistfile1.json

{
 "metadata": {
  "name": "",
  "signature": "sha256:4d284779bfca916c578238b9e9740b6a57238e9d3cd47386966ec7b613577c67"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Description\n",
      "\n",
      "Sometime we need to modify python for example to put all import * at the top of the code to don't get SyntaxWarning if the import are in the code (I know it's a dumb case).\n",
      "\n",
      "The first solution can be to use NodeTransformer from AST like that:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import ast\n",
      "\n",
      "class ImportParser(ast.NodeTransformer):\n",
      "\n",
      "    def visit_ImportFrom(self, node):\n",
      "        for a in node.names:\n",
      "            if a.name == '*':\n",
      "                ast.increment_lineno(node, -10)\n",
      "                node.col_offset = 0\n",
      "        return node"
     ],
     "language": "python",
     "metadata": {
      "slideshow": {
       "slide_type": "slide"
      }
     },
     "outputs": [],
     "prompt_number": 25
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "src = ast.parse('from sys import *\\n'\n",
      "                'def inner():\\n'\n",
      "                '    print \"test\"\\n'\n",
      "                '    from os import *\\n'\n",
      "                '    def innerinner():\\n'\n",
      "                '        from time import sleep\\n'\n",
      "                '        print \"inner\"\\n'\n",
      "                '    innerinner()\\n'\n",
      "                '    print \"coucou\"\\n'                \n",
      "                'inner()\\n')"
     ],
     "language": "python",
     "metadata": {
      "slideshow": {
       "slide_type": "slide"
      }
     },
     "outputs": [],
     "prompt_number": 26
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "src = ImportParser().visit(src)\n",
      "code = compile(src, '<bwah>', \"exec\", dont_inherit=True)\n",
      "\n",
      "ns = {}\n",
      "exec code in ns"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "test\n",
        "inner\n",
        "coucou\n"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stderr",
       "text": [
        "<bwah>:2: SyntaxWarning: import * only allowed at module level\n"
       ]
      }
     ],
     "prompt_number": 27
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Infortunately, it seems that the order of the instructions are not changed after that.\n",
      "So, we can try to change order manually like that:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from copy import deepcopy\n",
      "\n",
      "# A generator that return each instrution of the AST object\n",
      "def inst_generator(ast_tree):\n",
      "    yield ast_tree\n",
      "    if hasattr(ast_tree, 'body'):\n",
      "        for inst in ast_tree.body:\n",
      "            # can be replace by 'yield from inst_generator(inst) in python 3\n",
      "            for inst in inst_generator(inst):\n",
      "                yield inst\n",
      "\n",
      "# This method remove all import * from the tree\n",
      "def clear_tree(ast_tree):\n",
      "    if hasattr(ast_tree, 'body'):\n",
      "        for inst in list(ast_tree.body):\n",
      "            clear_tree(inst)\n",
      "            if (isinstance(inst, ast.ImportFrom) \n",
      "            and inst.names[0].name == '*'):\n",
      "                print vars(inst)\n",
      "                ast_tree.body.remove(inst)\n",
      "    return ast_tree\n",
      "                \n",
      "# This method copy the original tree, remove all import * from the new\n",
      "# and finally add each import * at the beginning of the tree\n",
      "def reorder_tree(ast_tree):\n",
      "    for inst in inst_generator(ast_tree):\n",
      "        if isinstance(inst, ast.Module):\n",
      "            tree_root = clear_tree(deepcopy(inst))\n",
      "        if (isinstance(inst, ast.ImportFrom) \n",
      "        and inst.names[0].name == '*'):\n",
      "            tree_root.body.insert(0, inst)\n",
      "    return tree_root\n",
      "\n",
      "new_src = reorder_tree(src)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "{'col_offset': 0, 'names': [<_ast.alias object at 0x308b9d0>], 'lineno': 1, 'module': 'sys', 'level': 0}\n",
        "{'col_offset': 4, 'names': [<_ast.alias object at 0x308bf90>], 'lineno': 4, 'module': 'os', 'level': 0}\n",
        "{'col_offset': 8, 'names': [<_ast.alias object at 0x308ba90>], 'lineno': 6, 'module': 'time', 'level': 0}\n"
       ]
      }
     ],
     "prompt_number": 24
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Now, we use unparse module (who can be found here: http://svn.python.org/projects/python/trunk/Demo/parser/unparse.py) to transform the code from AST to python.\n",
      "\n",
      "We can see that the import * are on the top of the code"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import unparse\n",
      "import sys\n",
      "unparse.Unparser(new_src)"
     ],
     "language": "python",
     "metadata": {
      "slideshow": {
       "slide_type": "slide"
      }
     },
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        "from os import *\n",
        "from sys import *\n",
        "\n",
        "def inner():\n",
        "    print 'test'\n",
        "\n",
        "    def innerinner():\n",
        "        from time import sleep\n",
        "        print 'inner'\n",
        "    innerinner()\n",
        "    print 'coucou'\n",
        "inner()"
       ]
      },
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 22,
       "text": [
        "<unparse.Unparser instance at 0x2f043f8>"
       ]
      }
     ],
     "prompt_number": 22
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Finally, we can execute the code without SyntaxWarning:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "code = compile(new_src, '<bwah>', \"exec\", dont_inherit=True)\n",
      "\n",
      "ns = {}\n",
      "exec code in ns"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "test\n",
        "inner\n",
        "coucou\n"
       ]
      }
     ],
     "prompt_number": 7
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def crawl_body(ast_obj, indent=0):\n",
      "    indent_str = '    '\n",
      "    print '%s%s' % (indent_str*indent, ast_obj)\n",
      "    if hasattr(ast_obj, 'body'):\n",
      "        for inst in ast_obj.body:\n",
      "            crawl_body(inst, indent=indent+1)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 15
    }
   ],
   "metadata": {}
  }
 ]
}