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  
       7  And this import:
       8      import spam
       9  Becomes:
      10      from . import spam
      11  """
      12  
      13  # Local imports
      14  from .. import fixer_base
      15  from os.path import dirname, join, exists, sep
      16  from ..fixer_util import FromImport, syms, token
      17  
      18  
      19  def traverse_imports(names):
      20      """
      21      Walks over all the names imported in a dotted_as_names node.
      22      """
      23      pending = [names]
      24      while pending:
      25          node = pending.pop()
      26          if node.type == token.NAME:
      27              yield node.value
      28          elif node.type == syms.dotted_name:
      29              yield "".join([ch.value for ch in node.children])
      30          elif node.type == syms.dotted_as_name:
      31              pending.append(node.children[0])
      32          elif node.type == syms.dotted_as_names:
      33              pending.extend(node.children[::-2])
      34          else:
      35              raise AssertionError("unknown node type")
      36  
      37  
      38  class ESC[4;38;5;81mFixImport(ESC[4;38;5;149mfixer_baseESC[4;38;5;149m.ESC[4;38;5;149mBaseFix):
      39      BM_compatible = True
      40  
      41      PATTERN = """
      42      import_from< 'from' imp=any 'import' ['('] any [')'] >
      43      |
      44      import_name< 'import' imp=any >
      45      """
      46  
      47      def start_tree(self, tree, name):
      48          super(FixImport, self).start_tree(tree, name)
      49          self.skip = "absolute_import" in tree.future_features
      50  
      51      def transform(self, node, results):
      52          if self.skip:
      53              return
      54          imp = results['imp']
      55  
      56          if node.type == syms.import_from:
      57              # Some imps are top-level (eg: 'import ham')
      58              # some are first level (eg: 'import ham.eggs')
      59              # some are third level (eg: 'import ham.eggs as spam')
      60              # Hence, the loop
      61              while not hasattr(imp, 'value'):
      62                  imp = imp.children[0]
      63              if self.probably_a_local_import(imp.value):
      64                  imp.value = "." + imp.value
      65                  imp.changed()
      66          else:
      67              have_local = False
      68              have_absolute = False
      69              for mod_name in traverse_imports(imp):
      70                  if self.probably_a_local_import(mod_name):
      71                      have_local = True
      72                  else:
      73                      have_absolute = True
      74              if have_absolute:
      75                  if have_local:
      76                      # We won't handle both sibling and absolute imports in the
      77                      # same statement at the moment.
      78                      self.warning(node, "absolute and local imports together")
      79                  return
      80  
      81              new = FromImport(".", [imp])
      82              new.prefix = node.prefix
      83              return new
      84  
      85      def probably_a_local_import(self, imp_name):
      86          if imp_name.startswith("."):
      87              # Relative imports are certainly not local imports.
      88              return False
      89          imp_name = imp_name.split(".", 1)[0]
      90          base_path = dirname(self.filename)
      91          base_path = join(base_path, imp_name)
      92          # If there is no __init__.py next to the file its not in a package
      93          # so can't be a relative import.
      94          if not exists(join(dirname(base_path), "__init__.py")):
      95              return False
      96          for ext in [".py", sep, ".pyc", ".so", ".sl", ".pyd"]:
      97              if exists(base_path + ext):
      98                  return True
      99          return False