cleanup-import-statement.py

#!/usr/bin/env python
"""Horrible hacked-together script to reformat import statements along
something that vaguely resembles our coding standard. To be called from
within Emacs via

(defun jvk/python-reformat-import ()
  (interactive)
  (save-excursion
    (let ((pmin (py-statement-point-begin))
          (pmax (py-statement-point-end)))
      (shell-command-on-region pmin pmax "<script.py>" nil nil))))

Properly speaking we should make this part of python.el in Emacs.

"""


import ast
import sys

LINE_WIDTH = 79
SEPARATOR = ', '
LINE_SEPARATOR = '\n' + 4 * ' '


def reformat_import_statement(module_name, import_names):
    import_names = sorted(import_names)
    from_statement = 'from {} import '.format(module_name)

    imports_line = SEPARATOR.join(import_names)
    if len(from_statement) + len(imports_line) < LINE_WIDTH:
        return from_statement + imports_line

    # need to split statement over several lines
    statement = from_statement + '(' + LINE_SEPARATOR
    current_length = 4
    for name in import_names[:-1]:
        if current_length + len(name) + len(SEPARATOR) > LINE_WIDTH:
            statement += LINE_SEPARATOR
            current_length = 4
        statement += name + SEPARATOR
        current_length += len(name) + len(SEPARATOR)
    if current_length + len(import_names[-1]) > LINE_WIDTH - 1:
        statement += LINE_SEPARATOR
    statement += import_names[-1] + '\n)'
    return statement


class ImportFromVisitor(ast.NodeVisitor):

    def visit_ImportFrom(self, node):
        module_name = node.module or '.'
        str_names = []
        for name_alias in node.names:
            s = name_alias.name
            if name_alias.asname:
                s += " as " + name_alias.asname
            str_names.append(s)
        print reformat_import_statement(module_name, str_names)
        self.generic_visit(node)


def cleanup_import_statement(stmt):
    tree = ast.parse(stmt)
    ImportFromVisitor().visit(tree)


if __name__ == '__main__':
    stmt = sys.stdin.read()
    cleanup_import_statement(stmt)