Merge pull request #242 from asottile/unsix-direct-imports-of-modules

Teach --py3-plus to rewrite direct six module imports

reorder_python_imports.py

@@ -244,6 +244,20 @@ def _inner() -> Generator[CodePartition, None, None]:
                         new_src = import_obj.to_text()
                         yield partition._replace(src=new_src)
                         break
+                    # from a.b.c import d => from c import d
+                    elif (
+                            (mod_parts + [symbol] == orig_mod) and
+                            not attr and
+                            len(new_mod) > 1 and
+                            (asname or symbol == new_mod[-1])
+                    ):
+                        import_obj.ast_obj.module = '.'.join(new_mod[:-1])
+                        import_obj.ast_obj.names = [
+                            ast.alias(name=new_mod[-1], asname=asname),
+                        ]
+                        new_src = import_obj.to_text()
+                        yield partition._replace(src=new_src)
+                        break
                     # from x.y import z => import z
                     elif (
                             not attr and

tests/reorder_python_imports_test.py

@@ -524,6 +524,54 @@ def test_replace_module_imported_asname():
     assert ret == 'import queue as Queue\n'
 
 
+def test_replace_module_imported_with_nested_replacement():
+    ret = fix_file_contents(
+        'from six.moves.urllib import parse\n',
+        imports_to_replace=[
+            (['six', 'moves', 'urllib', 'parse'], ['urllib', 'parse'], ''),
+        ],
+    )
+    assert ret == 'from urllib import parse\n'
+
+
+def test_replace_module_imported_with_nested_replacement_asname():
+    ret = fix_file_contents(
+        'from six.moves.urllib import parse as urllib_parse\n',
+        imports_to_replace=[
+            (['six', 'moves', 'urllib', 'parse'], ['urllib', 'parse'], ''),
+        ],
+    )
+    assert ret == 'from urllib import parse as urllib_parse\n'
+
+
+def test_replace_module_imported_with_nested_renamed_replacement_asname():
+    ret = fix_file_contents(
+        'from six.moves.urllib import parse as urllib_parse\n',
+        imports_to_replace=[
+            (['six', 'moves', 'urllib', 'parse'], ['urllib', 'parse2'], ''),
+        ],
+    )
+    assert ret == 'from urllib import parse2 as urllib_parse\n'
+
+
+def test_replace_module_skips_attr_specific_rules():
+    ret = fix_file_contents(
+        'from libone import util\n',
+        imports_to_replace=[
+            (['libone', 'util'], ['libtwo', 'util'], 'is_valid'),
+        ],
+    )
+    assert ret == 'from libone import util\n'
+
+
+def test_replace_module_skips_nonmatching_rules():
+    ret = fix_file_contents(
+        'from libthree import util\n',
+        imports_to_replace=[(['libone', 'util'], ['libtwo', 'util'], '')],
+    )
+    assert ret == 'from libthree import util\n'
+
+
 cases = pytest.mark.parametrize(
     ('s', 'expected'),
     (