1"""Fixer for import statements.
2If spam is being imported from the local directory, this import:
3    from spam import eggs
4Becomes:
5    from .spam import eggs
6
7And this import:
8    import spam
9Becomes:
10    from . import spam
11"""
12
13# Local imports
14from .. import fixer_base
15from os.path import dirname, join, exists, sep
16from ..fixer_util import FromImport, syms, token
17
18
19def 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("unkown node type")
36
37
38class FixImport(fixer_base.BaseFix):
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 = u"." + 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(u".", [imp])
82            new.prefix = node.prefix
83            return new
84
85    def probably_a_local_import(self, imp_name):
86        if imp_name.startswith(u"."):
87            # Relative imports are certainly not local imports.
88            return False
89        imp_name = imp_name.split(u".", 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
100