Lib/lib2to3/fixes/fix_imports.py

"""Fix incompatible imports and module references.

Fixes:
  * StringIO -> io
  * cStringIO -> io
  * md5 -> hashlib
"""
# Author: Collin Winter

# Local imports
from . import basefix
from .util import Name, attr_chain, any, set
import __builtin__
builtin_names = [name for name in dir(__builtin__)
                 if name not in ("__name__", "__doc__")]

MAPPING = {"StringIO":  ("io", ["StringIO"]),
           "cStringIO": ("io", ["StringIO"]),
           "__builtin__" : ("builtins", builtin_names), 
          }


def alternates(members):
    return "(" + "|".join(map(repr, members)) + ")"


def build_pattern():
    bare = set()
    for old_module, (new_module, members) in MAPPING.items():
        bare.add(old_module)
        bare.update(members)
        members = alternates(members)
        yield """import_name< 'import' (module=%r
                              | dotted_as_names< any* module=%r any* >) >
              """ % (old_module, old_module)
        yield """import_from< 'from' module_name=%r 'import'
                   ( %s | import_as_name< %s 'as' any >) >
              """ % (old_module, members, members)
        yield """import_from< 'from' module_name=%r 'import' star='*' >
              """ % old_module
        yield """import_name< 'import'
                              dotted_as_name< module_name=%r 'as' any > >
              """ % old_module
        yield """power< module_name=%r trailer< '.' %s > any* >
              """ % (old_module, members)
    yield """bare_name=%s""" % alternates(bare)


class FixImports(basefix.BaseFix):
    PATTERN = "|".join(build_pattern())

    order = "pre" # Pre-order tree traversal

    # Don't match the node if it's within another match
    def match(self, node):
        match = super(FixImports, self).match
        results = match(node)
        if results:
            if any([match(obj) for obj in attr_chain(node, "parent")]):
                return False
            return results
        return False

    def start_tree(self, tree, filename):
        super(FixImports, self).start_tree(tree, filename)
        self.replace = {}

    def transform(self, node, results):
        import_mod = results.get("module")
        mod_name = results.get("module_name")
        bare_name = results.get("bare_name")
        star = results.get("star")

        if import_mod or mod_name:
            new_name, members = MAPPING[(import_mod or mod_name).value]

        if import_mod:
            self.replace[import_mod.value] = new_name
            import_mod.replace(Name(new_name, prefix=import_mod.get_prefix()))
        elif mod_name:
            if star:
                self.cannot_convert(node, "Cannot handle star imports.")
            else:
                mod_name.replace(Name(new_name, prefix=mod_name.get_prefix()))
        elif bare_name:
            bare_name = bare_name[0]
            new_name = self.replace.get(bare_name.value)
            if new_name:
              bare_name.replace(Name(new_name, prefix=bare_name.get_prefix()))