~launchpad-pqm/launchpad/devel

#!/usr/bin/python
#
# Copyright 2010 Canonical Ltd.  This software is licensed under the
# GNU Affero General Public License version 3 (see the file LICENSE).

""" Format import sections in python files

= Usage =

format-imports <file or directory> ...

= Operation =

The script will process each filename on the command line. If the file is a
directory it recurses into it an process all *.py files found in the tree.
It will output the paths of all the files that have been changed.

For Launchpad it was applied to the "lib/canonical/launchpad" and the "lib/lp"
subtrees. Running it with those parameters on a freshly branched LP tree
should not produce any output, meaning that all the files in the tree should
be formatted correctly.

The script identifies the import section of each file as a block of lines
that start with "import" or "from" or are indented with at least one space or
are blank lines. Comment lines are also included if they are followed by an
import statement. An inital __future__ import and a module docstring are
explicitly skipped.

The import section is rewritten as three subsections, each separated by a
blank line. Any of the sections may be empty.
 1. Standard python library modules
 2. Import statements explicitly ordered to the top (see below)
 3. Third-party modules, meaning anything not fitting one of the other
    subsection criteria
 4. Local modules that begin with "canonical" or "lp".

Each section is sorted alphabetically by module name. Each module is put
on its own line, i.e.
{{{
  import os, sys
}}}
becomes
{{{
  import os
  import sys
}}}
Multiple import statements for the same module are conflated into one
statement, or two if the module was imported alongside an object inside it,
i.e.
{{{
  import sys
  from sys import stdin
}}}

Statements that import more than one objects are put on multiple lines in
list style, i.e.
{{{
  from sys import (
      stdin,
      stdout,
      )
}}}
Objects are sorted alphabetically and case-insensitively. One-object imports
are only formatted in this manner if the statement exceeds 78 characters in
length.

Comments stick with the import statement that followed them. Comments at the
end of one-line statements are moved to be be in front of it, .i.e.
{{{
  from sys import exit # Have a way out
}}}
becomes
{{{
  # Have a way out
  from sys import exit
}}}

= Format control =

Two special comments allow to control the operation of the formatter.

When an import statement is immediately preceded by a comment that starts
with the word "FIRST", it is placed into the second subsection (see above).

When the first import statement is directly preceded by a comment that starts
with the word "SKIP", the entire file is exempt from formatting.

= Known bugs =

Make sure to always check the result of the re-formatting to see if you have
been bitten by one of these.

Comments inside multi-line import statements break the formatter. A statement
like this will be ignored:
{{{
  from lp.app.interfaces import (
      # Don't do this.
      IMyInterface,
      IMyOtherInterface, # Don't do this either
      )
}}}
Actually, this will make the statement and all following to be ignored:
{{{
  from lp.app.interfaces import (
  # Breaks indentation rules anyway.
      IMyInterface,
      IMyOtherInterface,
      )
}}}

If a single-line statement has both a comment in front of it and at the end
of the line, only the end-line comment will survive. This could probably
easily be fixed to concatenate the too.
{{{
  # I am a gonner.
  from lp.app.interfaces import IMyInterface # I will survive!
}}}

Line continuation characters are recognized and resolved but
not re-introduced. This may leave the re-formatted text with a line that
is over the length limit.
{{{
    from lp.app.verylongnames.orverlydeep.modulestructure.leavenoroom \
        import object
}}}
"""

__metaclass__ = type

# SKIP this file when reformatting.
import os
import re
import sys
from textwrap import dedent

sys.path[0:0] = [os.path.dirname(__file__)]
from python_standard_libs import python_standard_libs

# python_standard_libs is only used for membership tests.
python_standard_libs = frozenset(python_standard_libs)

# To search for escaped newline chars.
escaped_nl_regex = re.compile("\\\\\n", re.M)
import_regex = re.compile("^import +(?P<module>.+)$", re.M)
from_import_single_regex = re.compile(
    "^from (?P<module>.+) +import +"
    "(?P<objects>[*]|[a-zA-Z0-9_, ]+)"
    "(?P<comment>#.*)?$", re.M)
from_import_multi_regex = re.compile(
    "^from +(?P<module>.+) +import *[(](?P<objects>[a-zA-Z0-9_, \n]+)[)]$",
    re.M)
comment_regex = re.compile(
    "(?P<comment>(^#.+\n)+)(^import|^from) +(?P<module>[a-zA-Z0-9_.]+)", re.M)
split_regex = re.compile(",\s*")

# Module docstrings are multiline (""") strings that are not indented and are
# followed at some point by an import .
module_docstring_regex = re.compile(
    '(?P<docstring>^["]{3}[^"]+["]{3}\n).*^(import |from .+ import)',
    re.M | re.S)
# The imports section starts with an import state that is not a __future__
# import and consists of import lines, indented lines, empty lines and
# comments which are followed by an import line. Sometimes we even find
# lines that contain a single ")"... :-(
imports_section_regex = re.compile(
    "(^#.+\n)*^(import|(from ((?!__future__)\S+) import)).*\n"
    "(^import .+\n|^from .+\n|^[\t ]+.+\n|(^#.+\n)+((^import|^from) "
    ".+\n)|^\n|^[)]\n)*",
    re.M)


def format_import_lines(module, objects):
    """Generate correct from...import strings."""
    if len(objects) == 1:
        statement = "from %s import %s" % (module, objects[0])
        if len(statement) < 79:
            return statement
    return "from %s import (\n    %s,\n    )" % (
        module, ",\n    ".join(objects))


def find_imports_section(content):
    """Return that part of the file that contains the import statements."""
    # Skip module docstring.
    match = module_docstring_regex.search(content)
    if match is None:
        startpos = 0
    else:
        startpos = match.end('docstring')

    match = imports_section_regex.search(content, startpos)
    if match is None:
        return (None, None)
    startpos = match.start()
    endpos = match.end()
    if content[startpos:endpos].startswith('# SKIP'):
        # Skip files explicitely.
        return(None, None)
    return (startpos, endpos)


class ImportStatement:
    """Holds information about an import statement."""

    def __init__(self, objects=None, comment=None):
        self.import_module = objects is None
        if objects is None:
            self.objects = None
        else:
            self.objects = sorted(objects, key=str.lower)
        self.comment = comment

    def addObjects(self, new_objects):
        """More objects in this statement; eliminate duplicates."""
        if self.objects is None:
            # No objects so far.
            self.objects = new_objects
        else:
            # Use set to eliminate double objects.
            more_objects = set(self.objects + new_objects)
            self.objects = sorted(list(more_objects), key=str.lower)

    def setComment(self, comment):
        """Add a comment to the statement."""
        self.comment = comment


def parse_import_statements(import_section):
    """Split the import section into statements.

    Returns a dictionary with the module as the key and the objects being
    imported as a sorted list of strings."""
    imports = {}
    # Search for escaped newlines and remove them.
    searchpos = 0
    while True:
        match = escaped_nl_regex.search(import_section, searchpos)
        if match is None:
            break
        start = match.start()
        end = match.end()
        import_section = import_section[:start] + import_section[end:]
        searchpos = start
    # Search for simple one-line import statements.
    searchpos = 0
    while True:
        match = import_regex.search(import_section, searchpos)
        if match is None:
            break
        # These imports are marked by a "None" value.
        # Multiple modules in one statement are split up.
        for module in split_regex.split(match.group('module').strip()):
            imports[module] = ImportStatement()
        searchpos = match.end()
    # Search for "from ... import" statements.
    for pattern in (from_import_single_regex, from_import_multi_regex):
        searchpos = 0
        while True:
            match = pattern.search(import_section, searchpos)
            if match is None:
                break
            import_objects = split_regex.split(
                match.group('objects').strip(" \n,"))
            module = match.group('module').strip()
            # Only one pattern has a 'comment' group.
            comment = match.groupdict().get('comment', None)
            if module in imports:
                # Catch double import lines.
                imports[module].addObjects(import_objects)
            else:
                imports[module] = ImportStatement(import_objects)
            if comment is not None:
                imports[module].setComment(comment)
            searchpos = match.end()
    # Search for comments in import section.
    searchpos = 0
    while True:
        match = comment_regex.search(import_section, searchpos)
        if match is None:
            break
        module = match.group('module').strip()
        comment = match.group('comment').strip()
        imports[module].setComment(comment)
        searchpos = match.end()

    return imports


def format_imports(imports):
    """Group and order imports, return the new import statements."""
    standard_section = {}
    first_section = {}
    thirdparty_section = {}
    local_section = {}
    # Group modules into sections.
    for module, statement in imports.iteritems():
        module_base = module.split('.')[0]
        comment = statement.comment
        if comment is not None and comment.startswith("# FIRST"):
            first_section[module] = statement
        elif module_base in ('canonical', 'lp'):
            local_section[module] = statement
        elif module_base in python_standard_libs:
            standard_section[module] = statement
        else:
            thirdparty_section[module] = statement

    all_import_lines = []
    # Sort within each section and generate statement strings.
    sections = (
        standard_section,
        first_section,
        thirdparty_section,
        local_section,
        )
    for section in sections:
        import_lines = []
        for module in sorted(section.keys(), key=str.lower):
            if section[module].comment is not None:
                import_lines.append(section[module].comment)
            if section[module].import_module:
                import_lines.append("import %s" % module)
            if section[module].objects is not None:
                import_lines.append(
                    format_import_lines(module, section[module].objects))
        if len(import_lines) > 0:
            all_import_lines.append('\n'.join(import_lines))
    # Sections are separated by two blank lines.
    return '\n\n'.join(all_import_lines)


def reformat_importsection(filename):
    """Replace the given file with a reformatted version of it."""
    pyfile = file(filename).read()
    import_start, import_end = find_imports_section(pyfile)
    if import_start is None:
        # Skip files with no import section.
        return False
    imports_section = pyfile[import_start:import_end]
    imports = parse_import_statements(imports_section)

    if pyfile[import_end:import_end + 1] != '#':
        # Two newlines before anything but comments.
        number_of_newlines = 3
    else:
        number_of_newlines = 2

    new_imports = format_imports(imports) + ("\n" * number_of_newlines)
    if new_imports == imports_section:
        # No change, no need to write a new file.
        return False

    new_file = open(filename, "w")
    new_file.write(pyfile[:import_start])
    new_file.write(new_imports)
    new_file.write(pyfile[import_end:])

    return True


def process_file(fpath):
    """Process the file with the given path."""
    changed = reformat_importsection(fpath)
    if changed:
        print fpath


def process_tree(dpath):
    """Walk a directory tree and process all *.py files."""
    for dirpath, dirnames, filenames in os.walk(dpath):
        for filename in filenames:
            if filename.endswith('.py'):
                process_file(os.path.join(dirpath, filename))


if __name__ == "__main__":
    if len(sys.argv) == 1 or sys.argv[1] in ("-h", "-?", "--help"):
        sys.stderr.write(dedent("""\
        usage: format-imports <file or directory> ...

        Type "format-imports --docstring | less" to see the documentation.
        """))
        sys.exit(1)
    if sys.argv[1] == "--docstring":
        sys.stdout.write(__doc__)
        sys.exit(2)
    for filename in sys.argv[1:]:
        if os.path.isdir(filename):
            process_tree(filename)
        else:
            process_file(filename)
    sys.exit(0)