cpython/Lib/lib2to3/fixes/fix_imports.py

90 lines
3.0 KiB
Python

"""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()))