From c43c7d95fe2e962df411360a088f3349dced0246 Mon Sep 17 00:00:00 2001
From: Aneesh Agrawal <aagrawal@lyft.com>
Date: Fri, 22 Apr 2022 10:49:49 -0400
Subject: [PATCH] Teach --py3-plus to rewrite direct six module imports

Impetus for this is to automatically rewrite
`from six.moves.urllib import parse` to `from urllib import parse`.
---
 reorder_python_imports.py            | 14 ++++++++
 tests/reorder_python_imports_test.py | 48 ++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/reorder_python_imports.py b/reorder_python_imports.py
index fc86f0e..445b31a 100644
--- a/reorder_python_imports.py
+++ b/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
diff --git a/tests/reorder_python_imports_test.py b/tests/reorder_python_imports_test.py
index 9fd49c8..a476762 100644
--- a/tests/reorder_python_imports_test.py
+++ b/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'),
     (