cpython/Lib/lib2to3/fixes/fix_import.py
Benjamin Peterson 0f6de936f8 Merged revisions 65141 via svnmerge from
svn+ssh://pythondev@svn.python.org/python/trunk

................
  r65141 | benjamin.peterson | 2008-07-19 09:14:06 -0500 (Sat, 19 Jul 2008) | 9 lines

  Merged revisions 65137 via svnmerge from
  svn+ssh://pythondev@svn.python.org/sandbox/trunk/2to3/lib2to3

  ........
    r65137 | georg.brandl | 2008-07-19 08:32:57 -0500 (Sat, 19 Jul 2008) | 2 lines

    #3334: correctly set prefix of imports.
  ........
................
2008-07-19 14:19:28 +00:00

60 lines
1.8 KiB
Python

"""Fixer for import statements.
If spam is being imported from the local directory, this import:
from spam import eggs
Becomes:
from .spam import eggs
And this import:
import spam
Becomes:
from . import spam
"""
# Local imports
from .. import fixer_base
from os.path import dirname, join, exists, pathsep
from ..fixer_util import FromImport
class FixImport(fixer_base.BaseFix):
PATTERN = """
import_from< type='from' imp=any 'import' any >
|
import_name< type='import' imp=any >
"""
def transform(self, node, results):
imp = results['imp']
if str(imp).startswith('.'):
# Already a new-style import
return
if not probably_a_local_import(str(imp), self.filename):
# I guess this is a global import -- skip it!
return
if results['type'].value == 'from':
# Some imps are top-level (eg: 'import ham')
# some are first level (eg: 'import ham.eggs')
# some are third level (eg: 'import ham.eggs as spam')
# Hence, the loop
while not hasattr(imp, 'value'):
imp = imp.children[0]
imp.value = "." + imp.value
node.changed()
else:
new = FromImport('.', getattr(imp, 'content', None) or [imp])
new.set_prefix(node.get_prefix())
node = new
return node
def probably_a_local_import(imp_name, file_path):
# Must be stripped because the right space is included by the parser
imp_name = imp_name.split('.', 1)[0].strip()
base_path = dirname(file_path)
base_path = join(base_path, imp_name)
for ext in ['.py', pathsep, '.pyc', '.so', '.sl', '.pyd']:
if exists(base_path + ext):
return True
return False