/Lib/lib2to3/fixes/fix_import.py

http://unladen-swallow.googlecode.com/ · Python · 90 lines · 57 code · 10 blank · 23 comment · 19 complexity · 9973617c9e868b2b9afb0a609ef30b35 MD5 · raw file

  1. """Fixer for import statements.
  2. If spam is being imported from the local directory, this import:
  3. from spam import eggs
  4. Becomes:
  5. from .spam import eggs
  6. And this import:
  7. import spam
  8. Becomes:
  9. from . import spam
  10. """
  11. # Local imports
  12. from .. import fixer_base
  13. from os.path import dirname, join, exists, pathsep
  14. from ..fixer_util import FromImport, syms, token
  15. def traverse_imports(names):
  16. """
  17. Walks over all the names imported in a dotted_as_names node.
  18. """
  19. pending = [names]
  20. while pending:
  21. node = pending.pop()
  22. if node.type == token.NAME:
  23. yield node.value
  24. elif node.type == syms.dotted_name:
  25. yield "".join([ch.value for ch in node.children])
  26. elif node.type == syms.dotted_as_name:
  27. pending.append(node.children[0])
  28. elif node.type == syms.dotted_as_names:
  29. pending.extend(node.children[::-2])
  30. else:
  31. raise AssertionError("unkown node type")
  32. class FixImport(fixer_base.BaseFix):
  33. PATTERN = """
  34. import_from< 'from' imp=any 'import' ['('] any [')'] >
  35. |
  36. import_name< 'import' imp=any >
  37. """
  38. def transform(self, node, results):
  39. imp = results['imp']
  40. if node.type == syms.import_from:
  41. # Some imps are top-level (eg: 'import ham')
  42. # some are first level (eg: 'import ham.eggs')
  43. # some are third level (eg: 'import ham.eggs as spam')
  44. # Hence, the loop
  45. while not hasattr(imp, 'value'):
  46. imp = imp.children[0]
  47. if self.probably_a_local_import(imp.value):
  48. imp.value = "." + imp.value
  49. imp.changed()
  50. return node
  51. else:
  52. have_local = False
  53. have_absolute = False
  54. for mod_name in traverse_imports(imp):
  55. if self.probably_a_local_import(mod_name):
  56. have_local = True
  57. else:
  58. have_absolute = True
  59. if have_absolute:
  60. if have_local:
  61. # We won't handle both sibling and absolute imports in the
  62. # same statement at the moment.
  63. self.warning(node, "absolute and local imports together")
  64. return
  65. new = FromImport('.', [imp])
  66. new.set_prefix(node.get_prefix())
  67. return new
  68. def probably_a_local_import(self, imp_name):
  69. imp_name = imp_name.split('.', 1)[0]
  70. base_path = dirname(self.filename)
  71. base_path = join(base_path, imp_name)
  72. # If there is no __init__.py next to the file its not in a package
  73. # so can't be a relative import.
  74. if not exists(join(dirname(base_path), '__init__.py')):
  75. return False
  76. for ext in ['.py', pathsep, '.pyc', '.so', '.sl', '.pyd']:
  77. if exists(base_path + ext):
  78. return True
  79. return False