1"""Fix changes imports of urllib which are now incompatible.
2   This is rather similar to fix_imports, but because of the more
3   complex nature of the fixing for urllib, it has its own fixer.
4"""
5# Author: Nick Edds
6
7# Local imports
8from lib2to3.fixes.fix_imports import alternates, FixImports
9from lib2to3.fixer_util import (Name, Comma, FromImport, Newline,
10                                find_indentation, Node, syms)
11
12MAPPING = {"urllib":  [
13                ("urllib.request",
14                    ["URLopener", "FancyURLopener", "urlretrieve",
15                     "_urlopener", "urlopen", "urlcleanup",
16                     "pathname2url", "url2pathname"]),
17                ("urllib.parse",
18                    ["quote", "quote_plus", "unquote", "unquote_plus",
19                     "urlencode", "splitattr", "splithost", "splitnport",
20                     "splitpasswd", "splitport", "splitquery", "splittag",
21                     "splittype", "splituser", "splitvalue", ]),
22                ("urllib.error",
23                    ["ContentTooShortError"])],
24           "urllib2" : [
25                ("urllib.request",
26                    ["urlopen", "install_opener", "build_opener",
27                     "Request", "OpenerDirector", "BaseHandler",
28                     "HTTPDefaultErrorHandler", "HTTPRedirectHandler",
29                     "HTTPCookieProcessor", "ProxyHandler",
30                     "HTTPPasswordMgr",
31                     "HTTPPasswordMgrWithDefaultRealm",
32                     "AbstractBasicAuthHandler",
33                     "HTTPBasicAuthHandler", "ProxyBasicAuthHandler",
34                     "AbstractDigestAuthHandler",
35                     "HTTPDigestAuthHandler", "ProxyDigestAuthHandler",
36                     "HTTPHandler", "HTTPSHandler", "FileHandler",
37                     "FTPHandler", "CacheFTPHandler",
38                     "UnknownHandler"]),
39                ("urllib.error",
40                    ["URLError", "HTTPError"]),
41           ]
42}
43
44# Duplicate the url parsing functions for urllib2.
45MAPPING["urllib2"].append(MAPPING["urllib"][1])
46
47
48def build_pattern():
49    bare = set()
50    for old_module, changes in MAPPING.items():
51        for change in changes:
52            new_module, members = change
53            members = alternates(members)
54            yield """import_name< 'import' (module=%r
55                                  | dotted_as_names< any* module=%r any* >) >
56                  """ % (old_module, old_module)
57            yield """import_from< 'from' mod_member=%r 'import'
58                       ( member=%s | import_as_name< member=%s 'as' any > |
59                         import_as_names< members=any*  >) >
60                  """ % (old_module, members, members)
61            yield """import_from< 'from' module_star=%r 'import' star='*' >
62                  """ % old_module
63            yield """import_name< 'import'
64                                  dotted_as_name< module_as=%r 'as' any > >
65                  """ % old_module
66            # bare_with_attr has a special significance for FixImports.match().
67            yield """power< bare_with_attr=%r trailer< '.' member=%s > any* >
68                  """ % (old_module, members)
69
70
71class FixUrllib(FixImports):
72
73    def build_pattern(self):
74        return "|".join(build_pattern())
75
76    def transform_import(self, node, results):
77        """Transform for the basic import case. Replaces the old
78           import name with a comma separated list of its
79           replacements.
80        """
81        import_mod = results.get("module")
82        pref = import_mod.prefix
83
84        names = []
85
86        # create a Node list of the replacement modules
87        for name in MAPPING[import_mod.value][:-1]:
88            names.extend([Name(name[0], prefix=pref), Comma()])
89        names.append(Name(MAPPING[import_mod.value][-1][0], prefix=pref))
90        import_mod.replace(names)
91
92    def transform_member(self, node, results):
93        """Transform for imports of specific module elements. Replaces
94           the module to be imported from with the appropriate new
95           module.
96        """
97        mod_member = results.get("mod_member")
98        pref = mod_member.prefix
99        member = results.get("member")
100
101        # Simple case with only a single member being imported
102        if member:
103            # this may be a list of length one, or just a node
104            if isinstance(member, list):
105                member = member[0]
106            new_name = None
107            for change in MAPPING[mod_member.value]:
108                if member.value in change[1]:
109                    new_name = change[0]
110                    break
111            if new_name:
112                mod_member.replace(Name(new_name, prefix=pref))
113            else:
114                self.cannot_convert(node, "This is an invalid module element")
115
116        # Multiple members being imported
117        else:
118            # a dictionary for replacements, order matters
119            modules = []
120            mod_dict = {}
121            members = results["members"]
122            for member in members:
123                # we only care about the actual members
124                if member.type == syms.import_as_name:
125                    as_name = member.children[2].value
126                    member_name = member.children[0].value
127                else:
128                    member_name = member.value
129                    as_name = None
130                if member_name != ",":
131                    for change in MAPPING[mod_member.value]:
132                        if member_name in change[1]:
133                            if change[0] not in mod_dict:
134                                modules.append(change[0])
135                            mod_dict.setdefault(change[0], []).append(member)
136
137            new_nodes = []
138            indentation = find_indentation(node)
139            first = True
140            def handle_name(name, prefix):
141                if name.type == syms.import_as_name:
142                    kids = [Name(name.children[0].value, prefix=prefix),
143                            name.children[1].clone(),
144                            name.children[2].clone()]
145                    return [Node(syms.import_as_name, kids)]
146                return [Name(name.value, prefix=prefix)]
147            for module in modules:
148                elts = mod_dict[module]
149                names = []
150                for elt in elts[:-1]:
151                    names.extend(handle_name(elt, pref))
152                    names.append(Comma())
153                names.extend(handle_name(elts[-1], pref))
154                new = FromImport(module, names)
155                if not first or node.parent.prefix.endswith(indentation):
156                    new.prefix = indentation
157                new_nodes.append(new)
158                first = False
159            if new_nodes:
160                nodes = []
161                for new_node in new_nodes[:-1]:
162                    nodes.extend([new_node, Newline()])
163                nodes.append(new_nodes[-1])
164                node.replace(nodes)
165            else:
166                self.cannot_convert(node, "All module elements are invalid")
167
168    def transform_dot(self, node, results):
169        """Transform for calls to module members in code."""
170        module_dot = results.get("bare_with_attr")
171        member = results.get("member")
172        new_name = None
173        if isinstance(member, list):
174            member = member[0]
175        for change in MAPPING[module_dot.value]:
176            if member.value in change[1]:
177                new_name = change[0]
178                break
179        if new_name:
180            module_dot.replace(Name(new_name,
181                                    prefix=module_dot.prefix))
182        else:
183            self.cannot_convert(node, "This is an invalid module element")
184
185    def transform(self, node, results):
186        if results.get("module"):
187            self.transform_import(node, results)
188        elif results.get("mod_member"):
189            self.transform_member(node, results)
190        elif results.get("bare_with_attr"):
191            self.transform_dot(node, results)
192        # Renaming and star imports are not supported for these modules.
193        elif results.get("module_star"):
194            self.cannot_convert(node, "Cannot handle star imports.")
195        elif results.get("module_as"):
196            self.cannot_convert(node, "This module is now multiple modules")
197