fix alias generation / units

Author: gdraheim

src/strip_python3.py

@@ -784,7 +784,7 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> Optional[ast.AST]:  # pylint
                     self.importfrom[modulename][symbol.name] = symbol.asname or symbol.name
                     origname = modulename + "." + symbol.name
                     codename = symbol.name if not symbol.asname else symbol.asname
-                    stmt = ast.ImportFrom(imports.module, [ast.alias(symbol.name, symbol.asname)], imports.level)
+                    stmt = ast.ImportFrom(imports.module, [ast.alias(symbol.name, symbol.asname if symbol.asname != symbol.name else None)], imports.level)
                     self.imported[origname] = stmt
                     self.asimport[codename] = origname
         return self.generic_visit(node)
@@ -793,7 +793,7 @@ def visit_Import(self, node: ast.Import) -> Optional[ast.AST]:  # pylint: disabl
         for symbol in imports.names:
             origname = symbol.name
             codename = symbol.name if not symbol.asname else symbol.asname
-            stmt = ast.Import([ast.alias(symbol.name, symbol.asname)])
+            stmt = ast.Import([ast.alias(symbol.name, symbol.asname if symbol.asname != symbol.name else None)])
             self.imported[origname] = stmt
             self.asimport[codename] = origname
         return self.generic_visit(node)
@@ -967,13 +967,12 @@ def visit(self, node: ast.AST) -> ast.AST:
                     body.append(stmt)
                 else:
                     if simple:
-                        body.append(ast.Import([ast.alias(mod, simple[mod]) for mod in sorted(simple)]))
+                        body.append(ast.Import([ast.alias(mod, simple[mod] if simple[mod] != mod else None) for mod in sorted(simple)]))
                     for mod in sorted(dotted):
                         alias = dotted[mod]
                         if alias and "." in mod:
                             libname, sym = mod.rsplit(".", 1)
-                            renamed = alias if sym != alias else None
-                            body.append(ast.ImportFrom(libname, [ast.alias(sym, renamed)], 0))
+                            body.append(ast.ImportFrom(libname, [ast.alias(sym, alias if alias != sym else None)], 0))
                         else:
                             body.append(ast.Import([ast.alias(mod, alias)]))
                     body.append(stmt)
@@ -989,13 +988,12 @@ def visit(self, node: ast.AST) -> ast.AST:
                     body.append(stmt)
                 else:
                     if simple:
-                        body.append(ast.Import([ast.alias(mod, simple[mod]) for mod in sorted(simple)]))
+                        body.append(ast.Import([ast.alias(mod, simple[mod] if simple[mod] != mod else None) for mod in sorted(simple)]))
                     for mod in sorted(dotted):
                         alias = dotted[mod]
                         if alias and "." in mod:
                             libname, sym = mod.rsplit(".", 1)
-                            renamed = alias if sym != alias else None
-                            body.append(ast.ImportFrom(libname, [ast.alias(sym, renamed)], 0))
+                            body.append(ast.ImportFrom(libname, [ast.alias(sym, alias if alias != sym else None)], 0))
                         else:
                             body.append(ast.Import([ast.alias(mod, alias)]))
                     body.append(stmt)

tests/unittests.py

@@ -7,11 +7,13 @@
 __author__ = "Guido U. Draheim"
 __version__ = "1.1.1127"
 
+from typing import cast
 import sys
 import unittest
 import logging
 import os.path
 from fnmatch import fnmatchcase as fnmatch
+import ast
 
 logg = logging.getLogger(os.path.basename(__file__))
 
@@ -90,6 +92,112 @@ def test_1112(self) -> None:
         self.assertEqual(c, "a\nb\n")
         self.assertEqual(d, "a\n        b\n")
         self.assertEqual(e, "a\n b\n")
+    def test_1201(self) -> None:
+        other = ast.Constant(1) # unknown element
+        have: ast.Module = app.pyi_module([other])
+        have0 = cast(ast.Constant, have.body[0])
+        self.assertEqual(other.value, have0.value)
+    def test_1210(self) -> None:
+        pyi = ast.parse(app.text4("""
+        def foo(a: A) -> B:
+            pass
+        """))
+        py1 = ast.parse(app.text4("""
+        from x import A
+        """))
+        py2 = ast.parse(app.text4("""
+        from y import B
+        """))
+        want = app.text4("""
+        from x import A
+        from y import B
+
+        def foo(a: A) -> B:
+            pass""")
+        pyi2 = app.pyi_copy_imports(pyi, py1, py2)
+        have = ast.unparse(pyi2) + "\n"
+        self.assertEqual(want, have)
+    def test_1211(self) -> None:
+        pyi = ast.parse(app.text4("""
+        def foo(a: A) -> B:
+            pass
+        """))
+        py1 = ast.parse(app.text4("""
+        from x.z import A
+        """))
+        py2 = ast.parse(app.text4("""
+        from y.z import B
+        """))
+        want = app.text4("""
+        from x.z import A
+        from y.z import B
+
+        def foo(a: A) -> B:
+            pass""")
+        pyi2 = app.pyi_copy_imports(pyi, py1, py2)
+        have = ast.unparse(pyi2) + "\n"
+        self.assertEqual(want, have)
+    def test_1212(self) -> None:
+        pyi = ast.parse(app.text4("""
+        def foo(a: x.A) -> y.B:
+            pass
+        """))
+        py1 = ast.parse(app.text4("""
+        import x
+        """))
+        py2 = ast.parse(app.text4("""
+        import y
+        """))
+        want = app.text4("""
+        import x, y
+
+        def foo(a: x.A) -> y.B:
+            pass""")
+        pyi2 = app.pyi_copy_imports(pyi, py1, py2)
+        have = ast.unparse(pyi2) + "\n"
+        self.assertEqual(want, have)
+    def test_1213(self) -> None:
+        pyi = ast.parse(app.text4("""
+        def foo(a: x.A) -> y.B:
+            pass
+        """))
+        py1 = ast.parse(app.text4("""
+        import app1.x as x
+        """))
+        py2 = ast.parse(app.text4("""
+        import app2.y as y
+        """))
+        want = app.text4("""
+        from app1 import x
+        from app2 import y
+
+        def foo(a: x.A) -> y.B:
+            pass""")
+        pyi2 = app.pyi_copy_imports(pyi, py1, py2)
+        have = ast.unparse(pyi2) + "\n"
+        self.assertEqual(want, have)
+    @unittest.expectedFailure
+    def test_1215(self) -> None:
+        pyi = ast.parse(app.text4("""
+        def foo(a: app1.x.A) -> app2.y.B:
+            pass
+        """))
+        py1 = ast.parse(app.text4("""
+        import app1.x
+        """))
+        py2 = ast.parse(app.text4("""
+        import app2.y
+        """))
+        want = app.text4("""
+        import app1.x
+        import app2.y
+
+        def foo(a: app1.x.A) -> app2.y.B:
+            pass""")
+        pyi2 = app.pyi_copy_imports(pyi, py1, py2)
+        have = ast.unparse(pyi2) + "\n"
+        self.assertEqual(want, have)
+
 
 if __name__ == "__main__":
     # unittest.main()