From 1d22a29d18afeda1eb878f3e905f91ddbe6af9df Mon Sep 17 00:00:00 2001
From: Zsolt Dollenstein <zsol.zsol@gmail.com>
Date: Thu, 1 Oct 2020 10:50:04 +0100
Subject: [PATCH] fix RemoveImportsVisitor crash when ImportAlias is inserted
 without comma (#397)

The comment preserving logic introduced in #392 assumed that in an ImportFrom node, ImportAliases have a comma property (except for the last one). That's only true if the ImportFrom node is parsed from actual source, but isn't necessarily true if it's constructed manually.
---
 libcst/codemod/visitors/_remove_imports.py    | 18 ++++++++----
 .../visitors/tests/test_remove_imports.py     | 28 ++++++++++++++++++-
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/libcst/codemod/visitors/_remove_imports.py b/libcst/codemod/visitors/_remove_imports.py
index 629fc021a..9d3b69023 100644
--- a/libcst/codemod/visitors/_remove_imports.py
+++ b/libcst/codemod/visitors/_remove_imports.py
@@ -374,12 +374,20 @@ def _process_importfrom_aliases(
                 if len(names_to_keep) != 0:
                     # there is a previous import alias
                     prev = names_to_keep[-1]
-                    names_to_keep[-1] = prev.with_deep_changes(
-                        whitespace_after=_merge_whitespace_after(
-                            prev.comma.whitespace_after,
-                            comma.whitespace_after,
+                    if isinstance(prev.comma, cst.Comma):
+                        prev = prev.with_deep_changes(
+                            prev.comma,
+                            whitespace_after=_merge_whitespace_after(
+                                prev.comma.whitespace_after,
+                                comma.whitespace_after,
+                            ),
                         )
-                    )
+                    else:
+                        # The previous alias didn't have a trailing comma. This can
+                        # occur if the alias was generated, instead of being parsed
+                        # from source.
+                        prev = prev.with_changes(comma=comma)
+                    names_to_keep[-1] = prev
                 else:
                     # No previous import alias, need to attach comment to `ImportFrom`.
                     # We can only do this if there was a leftparen on the import
diff --git a/libcst/codemod/visitors/tests/test_remove_imports.py b/libcst/codemod/visitors/tests/test_remove_imports.py
index f98aacbb6..564cf21e8 100644
--- a/libcst/codemod/visitors/tests/test_remove_imports.py
+++ b/libcst/codemod/visitors/tests/test_remove_imports.py
@@ -6,7 +6,7 @@
 import libcst as cst
 import libcst.matchers as m
 from libcst.codemod import CodemodContext, CodemodTest, VisitorBasedCodemodCommand
-from libcst.codemod.visitors import RemoveImportsVisitor
+from libcst.codemod.visitors import AddImportsVisitor, RemoveImportsVisitor
 from libcst.metadata import (
     QualifiedName,
     QualifiedNameProvider,
@@ -80,6 +80,12 @@ def test_remove_fromimport_keeping_standalone_comment(self) -> None:
                 short,
                 this_stays_too
             )
+            from fourth import (
+                a,
+                # comment
+                b,
+                c
+            )
         """
         after = """
             from foo import (
@@ -92,6 +98,10 @@ def test_remove_fromimport_keeping_standalone_comment(self) -> None:
             from third import (
                 this_stays_too
             )
+            from fourth import (
+                a,
+                c
+            )
         """
         self.assertCodemod(
             before,
@@ -101,6 +111,7 @@ def test_remove_fromimport_keeping_standalone_comment(self) -> None:
                 ("loooong", "short", None),
                 ("loooong", "bar", None),
                 ("third", "short", None),
+                ("fourth", "b", None),
             ],
         )
 
@@ -887,6 +898,21 @@ def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
             RemoveImportTransformer(CodemodContext()).transform_module(module).code,
         )
 
+    def test_remove_import_alias_after_inserting(self) -> None:
+        before = "from foo import bar, baz"
+        after = "from foo import quux, baz"
+
+        class AddRemoveTransformer(VisitorBasedCodemodCommand):
+            def visit_Module(self, node: cst.Module) -> None:
+                AddImportsVisitor.add_needed_import(self.context, "foo", "quux")
+                RemoveImportsVisitor.remove_unused_import(self.context, "foo", "bar")
+
+        module = cst.parse_module(self.make_fixture_data(before))
+        self.assertCodeEqual(
+            AddRemoveTransformer(CodemodContext()).transform_module(module).code,
+            after,
+        )
+
     def test_remove_comma(self) -> None:
         """
         Trailing commas should be removed if and only if the last alias is removed.