90 lines
3.0 KiB
Python
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()))
|