183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Copyright 2006 Google, Inc. All Rights Reserved.
283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Licensed to PSF under a Contributor Agreement.
383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh"""Refactoring framework.
583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew HsiehUsed as a main program, this can refactor any number of files and/or
783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehrecursively descend down directories.  Imported as a module, this
883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehprovides infrastructure to write your own refactoring tool.
983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh"""
1083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
1183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom __future__ import with_statement
1283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
1383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh__author__ = "Guido van Rossum <guido@python.org>"
1483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
1583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
1683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Python imports
1783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport os
1883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport sys
1983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport logging
2083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport operator
2183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport collections
2283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport StringIO
2383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom itertools import chain
2483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
2583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Local imports
2683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom .pgen2 import driver, tokenize, token
2783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom .fixer_util import find_root
2883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom . import pytree, pygram
2983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom . import btm_utils as bu
3083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom . import btm_matcher as bm
3183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
3283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
3383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef get_all_fix_names(fixer_pkg, remove_prefix=True):
3483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    """Return a sorted list of all available fix names in the given package."""
3583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    pkg = __import__(fixer_pkg, [], [], ["*"])
3683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    fixer_dir = os.path.dirname(pkg.__file__)
3783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    fix_names = []
3883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    for name in sorted(os.listdir(fixer_dir)):
3983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if name.startswith("fix_") and name.endswith(".py"):
4083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if remove_prefix:
4183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                name = name[4:]
4283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            fix_names.append(name[:-3])
4383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    return fix_names
4483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
4583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
4683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehclass _EveryNode(Exception):
4783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    pass
4883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
4983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
5083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef _get_head_types(pat):
5183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    """ Accepts a pytree Pattern Node and returns a set
5283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        of the pattern types which will match first. """
5383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
5483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if isinstance(pat, (pytree.NodePattern, pytree.LeafPattern)):
5583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # NodePatters must either have no type and no content
5683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        #   or a type and content -- so they don't get any farther
5783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # Always return leafs
5883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if pat.type is None:
5983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            raise _EveryNode
6083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return set([pat.type])
6183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
6283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if isinstance(pat, pytree.NegatedPattern):
6383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if pat.content:
6483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return _get_head_types(pat.content)
6583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        raise _EveryNode # Negated Patterns don't have a type
6683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
6783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if isinstance(pat, pytree.WildcardPattern):
6883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # Recurse on each node in content
6983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        r = set()
7083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for p in pat.content:
7183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for x in p:
7283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                r.update(_get_head_types(x))
7383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return r
7483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
7583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    raise Exception("Oh no! I don't understand pattern %s" %(pat))
7683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
7783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
7883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef _get_headnode_dict(fixer_list):
7983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    """ Accepts a list of fixers and returns a dictionary
8083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        of head node type --> fixer list.  """
8183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    head_nodes = collections.defaultdict(list)
8283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    every = []
8383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    for fixer in fixer_list:
8483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if fixer.pattern:
8583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            try:
8683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                heads = _get_head_types(fixer.pattern)
8783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            except _EveryNode:
8883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                every.append(fixer)
8983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
9083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                for node_type in heads:
9183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    head_nodes[node_type].append(fixer)
9283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
9383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if fixer._accept_type is not None:
9483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                head_nodes[fixer._accept_type].append(fixer)
9583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
9683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                every.append(fixer)
9783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    for node_type in chain(pygram.python_grammar.symbol2number.itervalues(),
9883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                           pygram.python_grammar.tokens):
9983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        head_nodes[node_type].extend(every)
10083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    return dict(head_nodes)
10183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
10283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
10383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef get_fixers_from_package(pkg_name):
10483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    """
10583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    Return the fully qualified names for fixers in the package pkg_name.
10683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    """
10783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    return [pkg_name + "." + fix_name
10883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for fix_name in get_all_fix_names(pkg_name, False)]
10983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
11083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef _identity(obj):
11183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    return obj
11283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
11383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehif sys.version_info < (3, 0):
11483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    import codecs
11583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    _open_with_encoding = codecs.open
11683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # codecs.open doesn't translate newlines sadly.
11783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def _from_system_newlines(input):
11883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return input.replace(u"\r\n", u"\n")
11983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def _to_system_newlines(input):
12083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if os.linesep != "\n":
12183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return input.replace(u"\n", os.linesep)
12283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
12383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return input
12483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehelse:
12583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    _open_with_encoding = open
12683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    _from_system_newlines = _identity
12783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    _to_system_newlines = _identity
12883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
12983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
13083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef _detect_future_features(source):
13183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    have_docstring = False
13283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    gen = tokenize.generate_tokens(StringIO.StringIO(source).readline)
13383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def advance():
13483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        tok = gen.next()
13583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return tok[0], tok[1]
13683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    ignore = frozenset((token.NEWLINE, tokenize.NL, token.COMMENT))
13783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    features = set()
13883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    try:
13983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        while True:
14083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            tp, value = advance()
14183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if tp in ignore:
14283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                continue
14383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            elif tp == token.STRING:
14483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if have_docstring:
14583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    break
14683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                have_docstring = True
14783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            elif tp == token.NAME and value == u"from":
14883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                tp, value = advance()
14983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if tp != token.NAME or value != u"__future__":
15083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    break
15183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                tp, value = advance()
15283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if tp != token.NAME or value != u"import":
15383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    break
15483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                tp, value = advance()
15583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if tp == token.OP and value == u"(":
15683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    tp, value = advance()
15783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                while tp == token.NAME:
15883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    features.add(value)
15983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    tp, value = advance()
16083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    if tp != token.OP or value != u",":
16183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        break
16283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    tp, value = advance()
16383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
16483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                break
16583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    except StopIteration:
16683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        pass
16783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    return frozenset(features)
16883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
16983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
17083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehclass FixerError(Exception):
17183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    """A fixer could not be loaded."""
17283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
17383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
17483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehclass RefactoringTool(object):
17583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
17683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    _default_options = {"print_function" : False,
17783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        "write_unchanged_files" : False}
17883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
17983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    CLASS_PREFIX = "Fix" # The prefix for fixer classes
18083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    FILE_PREFIX = "fix_" # The prefix for modules with a fixer within
18183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
18283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def __init__(self, fixer_names, options=None, explicit=None):
18383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Initializer.
18483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
18583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Args:
18683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            fixer_names: a list of fixers to import
18783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            options: an dict with configuration.
18883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            explicit: a list of fixers to run even if they are explicit.
18983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
19083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.fixers = fixer_names
19183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.explicit = explicit or []
19283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.options = self._default_options.copy()
19383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if options is not None:
19483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.options.update(options)
19583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if self.options["print_function"]:
19683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.grammar = pygram.python_grammar_no_print_statement
19783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
19883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.grammar = pygram.python_grammar
19983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # When this is True, the refactor*() methods will call write_file() for
20083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # files processed even if they were not changed during refactoring. If
20183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # and only if the refactor method's write parameter was True.
20283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.write_unchanged_files = self.options.get("write_unchanged_files")
20383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.errors = []
20483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.logger = logging.getLogger("RefactoringTool")
20583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.fixer_log = []
20683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.wrote = False
20783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.driver = driver.Driver(self.grammar,
20883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                    convert=pytree.convert,
20983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                    logger=self.logger)
21083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.pre_order, self.post_order = self.get_fixers()
21183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
21283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
21383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.files = []  # List of files that were or should be modified
21483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
21583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.BM = bm.BottomMatcher()
21683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.bmi_pre_order = [] # Bottom Matcher incompatible fixers
21783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.bmi_post_order = []
21883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
21983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for fixer in chain(self.post_order, self.pre_order):
22083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if fixer.BM_compatible:
22183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.BM.add_fixer(fixer)
22283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                # remove fixers that will be handled by the bottom-up
22383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                # matcher
22483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            elif fixer in self.pre_order:
22583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.bmi_pre_order.append(fixer)
22683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            elif fixer in self.post_order:
22783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.bmi_post_order.append(fixer)
22883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
22983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.bmi_pre_order_heads = _get_headnode_dict(self.bmi_pre_order)
23083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.bmi_post_order_heads = _get_headnode_dict(self.bmi_post_order)
23183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
23283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
23383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
23483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def get_fixers(self):
23583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Inspects the options to load the requested patterns and handlers.
23683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
23783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Returns:
23883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh          (pre_order, post_order), where pre_order is the list of fixers that
23983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh          want a pre-order AST traversal, and post_order is the list that want
24083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh          post-order traversal.
24183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
24283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        pre_order_fixers = []
24383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        post_order_fixers = []
24483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for fix_mod_path in self.fixers:
24583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            mod = __import__(fix_mod_path, {}, {}, ["*"])
24683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            fix_name = fix_mod_path.rsplit(".", 1)[-1]
24783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if fix_name.startswith(self.FILE_PREFIX):
24883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                fix_name = fix_name[len(self.FILE_PREFIX):]
24983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            parts = fix_name.split("_")
25083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            class_name = self.CLASS_PREFIX + "".join([p.title() for p in parts])
25183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            try:
25283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                fix_class = getattr(mod, class_name)
25383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            except AttributeError:
25483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                raise FixerError("Can't find %s.%s" % (fix_name, class_name))
25583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            fixer = fix_class(self.options, self.fixer_log)
25683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if fixer.explicit and self.explicit is not True and \
25783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    fix_mod_path not in self.explicit:
25883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.log_message("Skipping implicit fixer: %s", fix_name)
25983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                continue
26083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
26183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_debug("Adding transformation: %s", fix_name)
26283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if fixer.order == "pre":
26383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                pre_order_fixers.append(fixer)
26483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            elif fixer.order == "post":
26583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                post_order_fixers.append(fixer)
26683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
26783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                raise FixerError("Illegal fixer order: %r" % fixer.order)
26883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
26983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        key_func = operator.attrgetter("run_order")
27083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        pre_order_fixers.sort(key=key_func)
27183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        post_order_fixers.sort(key=key_func)
27283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return (pre_order_fixers, post_order_fixers)
27383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
27483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def log_error(self, msg, *args, **kwds):
27583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Called when an error occurs."""
27683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        raise
27783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
27883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def log_message(self, msg, *args):
27983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Hook to log a message."""
28083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if args:
28183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            msg = msg % args
28283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.logger.info(msg)
28383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
28483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def log_debug(self, msg, *args):
28583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if args:
28683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            msg = msg % args
28783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.logger.debug(msg)
28883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
28983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def print_output(self, old_text, new_text, filename, equal):
29083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Called with the old version, new version, and filename of a
29183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        refactored file."""
29283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        pass
29383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
29483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def refactor(self, items, write=False, doctests_only=False):
29583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Refactor a list of files and directories."""
29683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
29783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for dir_or_file in items:
29883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if os.path.isdir(dir_or_file):
29983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.refactor_dir(dir_or_file, write, doctests_only)
30083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
30183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.refactor_file(dir_or_file, write, doctests_only)
30283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
30383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def refactor_dir(self, dir_name, write=False, doctests_only=False):
30483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Descends down a directory and refactor every Python file found.
30583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
30683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Python files are assumed to have a .py extension.
30783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
30883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Files and subdirectories starting with '.' are skipped.
30983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
31083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        py_ext = os.extsep + "py"
31183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for dirpath, dirnames, filenames in os.walk(dir_name):
31283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_debug("Descending into %s", dirpath)
31383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            dirnames.sort()
31483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            filenames.sort()
31583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for name in filenames:
31683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if (not name.startswith(".") and
31783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    os.path.splitext(name)[1] == py_ext):
31883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    fullname = os.path.join(dirpath, name)
31983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    self.refactor_file(fullname, write, doctests_only)
32083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # Modify dirnames in-place to remove subdirs with leading dots
32183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            dirnames[:] = [dn for dn in dirnames if not dn.startswith(".")]
32283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
32383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def _read_python_source(self, filename):
32483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
32583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Do our best to decode a Python source file correctly.
32683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
32783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
32883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            f = open(filename, "rb")
32983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        except IOError as err:
33083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_error("Can't open %s: %s", filename, err)
33183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return None, None
33283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
33383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            encoding = tokenize.detect_encoding(f.readline)[0]
33483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        finally:
33583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            f.close()
33683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        with _open_with_encoding(filename, "r", encoding=encoding) as f:
33783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return _from_system_newlines(f.read()), encoding
33883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
33983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def refactor_file(self, filename, write=False, doctests_only=False):
34083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Refactors a file."""
34183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        input, encoding = self._read_python_source(filename)
34283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if input is None:
34383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # Reading the file failed.
34483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return
34583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        input += u"\n" # Silence certain parse errors
34683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if doctests_only:
34783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_debug("Refactoring doctests in %s", filename)
34883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            output = self.refactor_docstring(input, filename)
34983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if self.write_unchanged_files or output != input:
35083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.processed_file(output, filename, input, write, encoding)
35183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
35283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.log_debug("No doctest changes in %s", filename)
35383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
35483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            tree = self.refactor_string(input, filename)
35583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if self.write_unchanged_files or (tree and tree.was_changed):
35683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                # The [:-1] is to take off the \n we added earlier
35783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.processed_file(unicode(tree)[:-1], filename,
35883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                    write=write, encoding=encoding)
35983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
36083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.log_debug("No changes in %s", filename)
36183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
36283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def refactor_string(self, data, name):
36383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Refactor a given input string.
36483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
36583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Args:
36683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            data: a string holding the code to be refactored.
36783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            name: a human-readable name for use in error/log messages.
36883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
36983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Returns:
37083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            An AST corresponding to the refactored input stream; None if
37183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            there were errors during the parse.
37283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
37383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        features = _detect_future_features(data)
37483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if "print_function" in features:
37583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.driver.grammar = pygram.python_grammar_no_print_statement
37683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
37783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            tree = self.driver.parse_string(data)
37883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        except Exception as err:
37983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_error("Can't parse %s: %s: %s",
38083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                           name, err.__class__.__name__, err)
38183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return
38283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        finally:
38383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.driver.grammar = self.grammar
38483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        tree.future_features = features
38583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.log_debug("Refactoring %s", name)
38683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.refactor_tree(tree, name)
38783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return tree
38883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
38983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def refactor_stdin(self, doctests_only=False):
39083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        input = sys.stdin.read()
39183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if doctests_only:
39283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_debug("Refactoring doctests in stdin")
39383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            output = self.refactor_docstring(input, "<stdin>")
39483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if self.write_unchanged_files or output != input:
39583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.processed_file(output, "<stdin>", input)
39683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
39783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.log_debug("No doctest changes in stdin")
39883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
39983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            tree = self.refactor_string(input, "<stdin>")
40083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if self.write_unchanged_files or (tree and tree.was_changed):
40183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.processed_file(unicode(tree), "<stdin>", input)
40283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
40383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.log_debug("No changes in stdin")
40483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
40583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def refactor_tree(self, tree, name):
40683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Refactors a parse tree (modifying the tree in place).
40783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
40883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        For compatible patterns the bottom matcher module is
40983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        used. Otherwise the tree is traversed node-to-node for
41083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        matches.
41183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
41283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Args:
41383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            tree: a pytree.Node instance representing the root of the tree
41483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                  to be refactored.
41583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            name: a human-readable name for this tree.
41683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
41783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Returns:
41883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            True if the tree was modified, False otherwise.
41983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
42083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
42183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for fixer in chain(self.pre_order, self.post_order):
42283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            fixer.start_tree(tree, name)
42383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
42483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        #use traditional matching for the incompatible fixers
42583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.traverse_by(self.bmi_pre_order_heads, tree.pre_order())
42683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.traverse_by(self.bmi_post_order_heads, tree.post_order())
42783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
42883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # obtain a set of candidate nodes
42983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        match_set = self.BM.run(tree.leaves())
43083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
43183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        while any(match_set.values()):
43283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for fixer in self.BM.fixers:
43383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if fixer in match_set and match_set[fixer]:
43483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    #sort by depth; apply fixers from bottom(of the AST) to top
43583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    match_set[fixer].sort(key=pytree.Base.depth, reverse=True)
43683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
43783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    if fixer.keep_line_order:
43883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        #some fixers(eg fix_imports) must be applied
43983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        #with the original file's line order
44083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        match_set[fixer].sort(key=pytree.Base.get_lineno)
44183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
44283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    for node in list(match_set[fixer]):
44383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        if node in match_set[fixer]:
44483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            match_set[fixer].remove(node)
44583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
44683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        try:
44783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            find_root(node)
44883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        except ValueError:
44983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            # this node has been cut off from a
45083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            # previous transformation ; skip
45183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            continue
45283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
45383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        if node.fixers_applied and fixer in node.fixers_applied:
45483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            # do not apply the same fixer again
45583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            continue
45683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
45783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        results = fixer.match(node)
45883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
45983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        if results:
46083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            new = fixer.transform(node, results)
46183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            if new is not None:
46283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                node.replace(new)
46383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                #new.fixers_applied.append(fixer)
46483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                for node in new.post_order():
46583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                    # do not apply the fixer again to
46683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                    # this or any subnode
46783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                    if not node.fixers_applied:
46883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                        node.fixers_applied = []
46983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                    node.fixers_applied.append(fixer)
47083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
47183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                # update the original match set for
47283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                # the added code
47383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                new_matches = self.BM.run(new.leaves())
47483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                for fxr in new_matches:
47583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                    if not fxr in match_set:
47683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                        match_set[fxr]=[]
47783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
47883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                    match_set[fxr].extend(new_matches[fxr])
47983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
48083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for fixer in chain(self.pre_order, self.post_order):
48183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            fixer.finish_tree(tree, name)
48283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return tree.was_changed
48383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
48483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def traverse_by(self, fixers, traversal):
48583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Traverse an AST, applying a set of fixers to each node.
48683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
48783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        This is a helper method for refactor_tree().
48883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
48983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Args:
49083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            fixers: a list of fixer instances.
49183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            traversal: a generator that yields AST nodes.
49283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
49383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Returns:
49483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            None
49583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
49683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if not fixers:
49783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return
49883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for node in traversal:
49983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for fixer in fixers[node.type]:
50083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                results = fixer.match(node)
50183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if results:
50283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    new = fixer.transform(node, results)
50383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    if new is not None:
50483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        node.replace(new)
50583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        node = new
50683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
50783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def processed_file(self, new_text, filename, old_text=None, write=False,
50883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                       encoding=None):
50983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
51083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Called when a file has been refactored and there may be changes.
51183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
51283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.files.append(filename)
51383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if old_text is None:
51483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            old_text = self._read_python_source(filename)[0]
51583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if old_text is None:
51683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                return
51783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        equal = old_text == new_text
51883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.print_output(old_text, new_text, filename, equal)
51983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if equal:
52083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_debug("No changes to %s", filename)
52183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if not self.write_unchanged_files:
52283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                return
52383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if write:
52483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.write_file(new_text, filename, old_text, encoding)
52583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
52683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_debug("Not writing changes to %s", filename)
52783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
52883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def write_file(self, new_text, filename, old_text, encoding=None):
52983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Writes a string to a file.
53083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
53183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        It first shows a unified diff between the old text and the new text, and
53283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        then rewrites the file; the latter is only done if the write option is
53383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        set.
53483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
53583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
53683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            f = _open_with_encoding(filename, "w", encoding=encoding)
53783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        except os.error as err:
53883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_error("Can't create %s: %s", filename, err)
53983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return
54083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
54183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            f.write(_to_system_newlines(new_text))
54283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        except os.error as err:
54383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_error("Can't write %s: %s", filename, err)
54483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        finally:
54583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            f.close()
54683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.log_debug("Wrote changes to %s", filename)
54783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.wrote = True
54883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
54983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    PS1 = ">>> "
55083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    PS2 = "... "
55183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
55283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def refactor_docstring(self, input, filename):
55383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Refactors a docstring, looking for doctests.
55483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
55583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        This returns a modified version of the input string.  It looks
55683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for doctests, which start with a ">>>" prompt, and may be
55783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        continued with "..." prompts, as long as the "..." is indented
55883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        the same as the ">>>".
55983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
56083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        (Unfortunately we can't use the doctest module's parser,
56183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        since, like most parsers, it is not geared towards preserving
56283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        the original source.)
56383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
56483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        result = []
56583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        block = None
56683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        block_lineno = None
56783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        indent = None
56883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        lineno = 0
56983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for line in input.splitlines(True):
57083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            lineno += 1
57183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if line.lstrip().startswith(self.PS1):
57283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if block is not None:
57383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    result.extend(self.refactor_doctest(block, block_lineno,
57483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                                        indent, filename))
57583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                block_lineno = lineno
57683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                block = [line]
57783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                i = line.find(self.PS1)
57883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                indent = line[:i]
57983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            elif (indent is not None and
58083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                  (line.startswith(indent + self.PS2) or
58183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                   line == indent + self.PS2.rstrip() + u"\n")):
58283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                block.append(line)
58383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
58483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if block is not None:
58583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    result.extend(self.refactor_doctest(block, block_lineno,
58683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                                        indent, filename))
58783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                block = None
58883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                indent = None
58983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                result.append(line)
59083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if block is not None:
59183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            result.extend(self.refactor_doctest(block, block_lineno,
59283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                                indent, filename))
59383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return u"".join(result)
59483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
59583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def refactor_doctest(self, block, lineno, indent, filename):
59683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Refactors one doctest.
59783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
59883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        A doctest is given as a block of lines, the first of which starts
59983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        with ">>>" (possibly indented), while the remaining lines start
60083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        with "..." (identically indented).
60183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
60283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
60383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
60483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            tree = self.parse_block(block, lineno, indent)
60583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        except Exception as err:
60683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if self.logger.isEnabledFor(logging.DEBUG):
60783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                for line in block:
60883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    self.log_debug("Source: %s", line.rstrip(u"\n"))
60983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_error("Can't parse docstring in %s line %s: %s: %s",
61083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                           filename, lineno, err.__class__.__name__, err)
61183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return block
61283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if self.refactor_tree(tree, filename):
61383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            new = unicode(tree).splitlines(True)
61483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # Undo the adjustment of the line numbers in wrap_toks() below.
61583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            clipped, new = new[:lineno-1], new[lineno-1:]
61683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            assert clipped == [u"\n"] * (lineno-1), clipped
61783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if not new[-1].endswith(u"\n"):
61883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                new[-1] += u"\n"
61983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            block = [indent + self.PS1 + new.pop(0)]
62083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if new:
62183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                block += [indent + self.PS2 + line for line in new]
62283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return block
62383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
62483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def summarize(self):
62583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if self.wrote:
62683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            were = "were"
62783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
62883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            were = "need to be"
62983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if not self.files:
63083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_message("No files %s modified.", were)
63183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
63283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_message("Files that %s modified:", were)
63383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for file in self.files:
63483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.log_message(file)
63583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if self.fixer_log:
63683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.log_message("Warnings/messages while refactoring:")
63783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for message in self.fixer_log:
63883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.log_message(message)
63983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if self.errors:
64083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if len(self.errors) == 1:
64183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.log_message("There was 1 error:")
64283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
64383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.log_message("There were %d errors:", len(self.errors))
64483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for msg, args, kwds in self.errors:
64583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.log_message(msg, *args, **kwds)
64683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
64783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def parse_block(self, block, lineno, indent):
64883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Parses a block into a tree.
64983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
65083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        This is necessary to get correct line number / offset information
65183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        in the parser diagnostics and embedded into the parse tree.
65283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
65383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        tree = self.driver.parse_tokens(self.wrap_toks(block, lineno, indent))
65483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        tree.future_features = frozenset()
65583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return tree
65683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
65783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def wrap_toks(self, block, lineno, indent):
65883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Wraps a tokenize stream to systematically modify start/end."""
65983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        tokens = tokenize.generate_tokens(self.gen_lines(block, indent).next)
66083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for type, value, (line0, col0), (line1, col1), line_text in tokens:
66183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            line0 += lineno - 1
66283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            line1 += lineno - 1
66383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # Don't bother updating the columns; this is too complicated
66483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # since line_text would also have to be updated and it would
66583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # still break for tokens spanning lines.  Let the user guess
66683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # that the column numbers for doctests are relative to the
66783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # end of the prompt string (PS1 or PS2).
66883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            yield type, value, (line0, col0), (line1, col1), line_text
66983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
67083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
67183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def gen_lines(self, block, indent):
67283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """Generates lines as expected by tokenize from a list of lines.
67383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
67483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        This strips the first len(indent + self.PS1) characters off each line.
67583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        """
67683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        prefix1 = indent + self.PS1
67783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        prefix2 = indent + self.PS2
67883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        prefix = prefix1
67983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for line in block:
68083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if line.startswith(prefix):
68183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                yield line[len(prefix):]
68283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            elif line == prefix.rstrip() + u"\n":
68383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                yield u"\n"
68483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
68583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                raise AssertionError("line=%r, prefix=%r" % (line, prefix))
68683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            prefix = prefix2
68783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        while True:
68883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            yield ""
68983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
69083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
69183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehclass MultiprocessingUnsupported(Exception):
69283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    pass
69383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
69483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
69583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehclass MultiprocessRefactoringTool(RefactoringTool):
69683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
69783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def __init__(self, *args, **kwargs):
69883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        super(MultiprocessRefactoringTool, self).__init__(*args, **kwargs)
69983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.queue = None
70083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.output_lock = None
70183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
70283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def refactor(self, items, write=False, doctests_only=False,
70383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                 num_processes=1):
70483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if num_processes == 1:
70583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return super(MultiprocessRefactoringTool, self).refactor(
70683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                items, write, doctests_only)
70783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
70883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            import multiprocessing
70983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        except ImportError:
71083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            raise MultiprocessingUnsupported
71183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if self.queue is not None:
71283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            raise RuntimeError("already doing multiple processes")
71383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.queue = multiprocessing.JoinableQueue()
71483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self.output_lock = multiprocessing.Lock()
71583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        processes = [multiprocessing.Process(target=self._child)
71683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                     for i in xrange(num_processes)]
71783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
71883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for p in processes:
71983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                p.start()
72083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            super(MultiprocessRefactoringTool, self).refactor(items, write,
72183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                                              doctests_only)
72283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        finally:
72383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.queue.join()
72483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for i in xrange(num_processes):
72583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.queue.put(None)
72683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for p in processes:
72783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if p.is_alive():
72883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    p.join()
72983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.queue = None
73083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
73183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def _child(self):
73283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        task = self.queue.get()
73383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        while task is not None:
73483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            args, kwargs = task
73583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            try:
73683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                super(MultiprocessRefactoringTool, self).refactor_file(
73783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    *args, **kwargs)
73883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            finally:
73983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                self.queue.task_done()
74083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            task = self.queue.get()
74183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
74283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def refactor_file(self, *args, **kwargs):
74383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if self.queue is not None:
74483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.queue.put((args, kwargs))
74583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
74683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return super(MultiprocessRefactoringTool, self).refactor_file(
74783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                *args, **kwargs)
748