fix_import.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  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, sep
  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("unknown node type")
  32. class FixImport(fixer_base.BaseFix):
  33. BM_compatible = True
  34. PATTERN = """
  35. import_from< 'from' imp=any 'import' ['('] any [')'] >
  36. |
  37. import_name< 'import' imp=any >
  38. """
  39. def start_tree(self, tree, name):
  40. super(FixImport, self).start_tree(tree, name)
  41. self.skip = "absolute_import" in tree.future_features
  42. def transform(self, node, results):
  43. if self.skip:
  44. return
  45. imp = results['imp']
  46. if node.type == syms.import_from:
  47. # Some imps are top-level (eg: 'import ham')
  48. # some are first level (eg: 'import ham.eggs')
  49. # some are third level (eg: 'import ham.eggs as spam')
  50. # Hence, the loop
  51. while not hasattr(imp, 'value'):
  52. imp = imp.children[0]
  53. if self.probably_a_local_import(imp.value):
  54. imp.value = "." + imp.value
  55. imp.changed()
  56. else:
  57. have_local = False
  58. have_absolute = False
  59. for mod_name in traverse_imports(imp):
  60. if self.probably_a_local_import(mod_name):
  61. have_local = True
  62. else:
  63. have_absolute = True
  64. if have_absolute:
  65. if have_local:
  66. # We won't handle both sibling and absolute imports in the
  67. # same statement at the moment.
  68. self.warning(node, "absolute and local imports together")
  69. return
  70. new = FromImport(".", [imp])
  71. new.prefix = node.prefix
  72. return new
  73. def probably_a_local_import(self, imp_name):
  74. if imp_name.startswith("."):
  75. # Relative imports are certainly not local imports.
  76. return False
  77. imp_name = imp_name.split(".", 1)[0]
  78. base_path = dirname(self.filename)
  79. base_path = join(base_path, imp_name)
  80. # If there is no __init__.py next to the file its not in a package
  81. # so can't be a relative import.
  82. if not exists(join(dirname(base_path), "__init__.py")):
  83. return False
  84. for ext in [".py", sep, ".pyc", ".so", ".sl", ".pyd"]:
  85. if exists(base_path + ext):
  86. return True
  87. return False