fix_absolute_import.py 3.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. """
  2. Fixer for import statements, with a __future__ import line.
  3. Based on lib2to3/fixes/fix_import.py, but extended slightly so it also
  4. supports Cython modules.
  5. If spam is being imported from the local directory, this import:
  6. from spam import eggs
  7. becomes:
  8. from __future__ import absolute_import
  9. from .spam import eggs
  10. and this import:
  11. import spam
  12. becomes:
  13. from __future__ import absolute_import
  14. from . import spam
  15. """
  16. from os.path import dirname, join, exists, sep
  17. from lib2to3.fixes.fix_import import FixImport
  18. from lib2to3.fixer_util import FromImport, syms
  19. from lib2to3.fixes.fix_import import traverse_imports
  20. from libfuturize.fixer_util import future_import
  21. class FixAbsoluteImport(FixImport):
  22. run_order = 9
  23. def transform(self, node, results):
  24. """
  25. Copied from FixImport.transform(), but with this line added in
  26. any modules that had implicit relative imports changed:
  27. from __future__ import absolute_import"
  28. """
  29. if self.skip:
  30. return
  31. imp = results['imp']
  32. if node.type == syms.import_from:
  33. # Some imps are top-level (eg: 'import ham')
  34. # some are first level (eg: 'import ham.eggs')
  35. # some are third level (eg: 'import ham.eggs as spam')
  36. # Hence, the loop
  37. while not hasattr(imp, 'value'):
  38. imp = imp.children[0]
  39. if self.probably_a_local_import(imp.value):
  40. imp.value = u"." + imp.value
  41. imp.changed()
  42. future_import(u"absolute_import", node)
  43. else:
  44. have_local = False
  45. have_absolute = False
  46. for mod_name in traverse_imports(imp):
  47. if self.probably_a_local_import(mod_name):
  48. have_local = True
  49. else:
  50. have_absolute = True
  51. if have_absolute:
  52. if have_local:
  53. # We won't handle both sibling and absolute imports in the
  54. # same statement at the moment.
  55. self.warning(node, "absolute and local imports together")
  56. return
  57. new = FromImport(u".", [imp])
  58. new.prefix = node.prefix
  59. future_import(u"absolute_import", node)
  60. return new
  61. def probably_a_local_import(self, imp_name):
  62. """
  63. Like the corresponding method in the base class, but this also
  64. supports Cython modules.
  65. """
  66. if imp_name.startswith(u"."):
  67. # Relative imports are certainly not local imports.
  68. return False
  69. imp_name = imp_name.split(u".", 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", sep, ".pyc", ".so", ".sl", ".pyd", ".pyx"]:
  77. if exists(base_path + ext):
  78. return True
  79. return False