text: # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  #
  import libcst as cst
  import libcst.matchers as m
  from libcst.codemod import CodemodContext, CodemodTest, VisitorBasedCodemodCommand
  from libcst.codemod.visitors import AddImportsVisitor, RemoveImportsVisitor
  from libcst.metadata import (
      QualifiedName,
      QualifiedNameProvider,
      QualifiedNameSource,
      ScopeProvider,
  )
  from libcst.testing.utils import data_provider
  
  
  class TestRemoveImportsCodemod(CodemodTest):
  
      TRANSFORM = RemoveImportsVisitor
  
      def test_noop(self) -> None:
          """
          Should do nothing.
          """
  
          before = """
              def foo() -> None:
                  pass
          """
          after = """
              def foo() -> None:
                  pass
          """
  
          self.assertCodemod(before, after, [])
  
      def test_remove_import_simple(self) -> None:
          """
          Should remove module as import
          """
  
          before = """
              import bar
              import baz
  
              def foo() -> None:
                  pass
          """
          after = """
              import bar
  
              def foo() -> None:
                  pass
          """
  
          self.assertCodemod(before, after, [("baz", None, None)])
  
      def test_remove_fromimport_simple(self) -> None:
          before = "from a import b, c"
          after = "from a import c"
          self.assertCodemod(before, after, [("a", "b", None)])
  
      def test_remove_fromimport_keeping_standalone_comment(self) -> None:
          before = """
              from foo import (
                  bar,
                  # comment
                  baz,
              )
              from loooong import (
                  bar,
                  # comment
                  short,
                  this_stays
              )
              from third import (
                  # comment
                  short,
                  this_stays_too
              )
              from fourth import (
                  a,
                  # comment
                  b,
                  c
              )
          """
          after = """
              from foo import (
                  # comment
                  baz,
              )
              from loooong import (
                  this_stays
              )
              from third import (
                  this_stays_too
              )
              from fourth import (
                  a,
                  c
              )
          """
          self.assertCodemod(
              before,
              after,
              [
                  ("foo", "bar", None),
                  ("loooong", "short", None),
                  ("loooong", "bar", None),
                  ("third", "short", None),
                  ("fourth", "b", None),
              ],
          )
  
      def test_remove_fromimport_keeping_inline_comment(self) -> None:
          before = """
              from foo import (  # comment
                  bar,
                  # comment2
                  baz,
              )
              from loooong import (
                  bar,
                  short,  # comment
                  # comment2
                  this_stays
              )
              from third import (
                  short,  # comment
                  this_stays_too  # comment2
              )
          """
          after = """
              from foo import (  # comment
                  # comment2
                  baz,
              )
              from loooong import (
                  # comment2
                  this_stays
              )
              from third import (
                  this_stays_too  # comment2
              )
          """
          self.assertCodemod(
              before,
              after,
              [
                  ("foo", "bar", None),
                  ("loooong", "short", None),
                  ("loooong", "bar", None),
                  ("third", "short", None),
              ],
          )
  
      def test_remove_import_alias_simple(self) -> None:
          """
          Should remove aliased module as import
          """
  
          before = """
              import bar
              import baz as qux
  
              def foo() -> None:
                  pass
          """
          after = """
              import bar
  
              def foo() -> None:
                  pass
          """
  
          self.assertCodemod(before, after, [("baz", None, "qux")])
  
      def test_dont_remove_import_simple(self) -> None:
          """
          Should not remove module import with reference
          """
  
          before = """
              import bar
              import baz
  
              def foo() -> None:
                  baz.qux()
          """
          after = """
              import bar
              import baz
  
              def foo() -> None:
                  baz.qux()
          """
  
          self.assertCodemod(before, after, [("baz", None, None)])
  
      def test_dont_remove_import_alias_simple(self) -> None:
          """
          Should not remove aliased module import with reference
          """
  
          before = """
              import bar
              import baz as qux
  
              def foo() -> None:
                  qux.quux()
          """
          after = """
              import bar
              import baz as qux
  
              def foo() -> None:
                  qux.quux()
          """
  
          self.assertCodemod(before, after, [("baz", None, "qux")])
  
      def test_dont_remove_import_simple_wrong_alias(self) -> None:
          """
          Should not remove module as import since wrong alias
          """
  
          before = """
              import bar
              import baz
  
              def foo() -> None:
                  pass
          """
          after = """
              import bar
              import baz
  
              def foo() -> None:
                  pass
          """
  
          self.assertCodemod(before, after, [("baz", None, "qux")])
  
      def test_dont_remove_import_wrong_alias_simple(self) -> None:
          """
          Should not remove wrong aliased module as import
          """
  
          before = """
              import bar
              import baz as qux
  
              def foo() -> None:
                  pass
          """
          after = """
              import bar
              import baz as qux
  
              def foo() -> None:
                  pass
          """
  
          self.assertCodemod(before, after, [("baz", None, None)])
  
      def test_remove_importfrom_simple(self) -> None:
          """
          Should remove import from
          """
  
          before = """
              import bar
              from baz import qux
  
              def foo() -> None:
                  pass
          """
          after = """
              import bar
  
              def foo() -> None:
                  pass
          """
  
          self.assertCodemod(before, after, [("baz", "qux", None)])
  
      def test_remove_importfrom_alias_simple(self) -> None:
          """
          Should remove import from with alias
          """
  
          before = """
              import bar
              from baz import qux as quux
  
              def foo() -> None:
                  pass
          """
          after = """
              import bar
  
              def foo() -> None:
                  pass
          """
  
          self.assertCodemod(before, after, [("baz", "qux", "quux")])
  
      def test_dont_remove_importfrom_simple(self) -> None:
          """
          Should not remove import from with reference
          """
  
          before = """
              import bar
              from baz import qux
  
              def foo() -> None:
                  qux()
          """
          after = """
              import bar
              from baz import qux
  
              def foo() -> None:
                  qux()
          """
  
          self.assertCodemod(before, after, [("baz", "qux", None)])
  
      def test_dont_remove_importfrom_alias_simple(self) -> None:
          """
          Should not remove aliased import from with reference
          """
  
          before = """
              import bar
              from baz import qux as quux
  
              def foo() -> None:
                  quux()
          """
          after = """
              import bar
              from baz import qux as quux
  
              def foo() -> None:
                  quux()
          """
  
          self.assertCodemod(before, after, [("baz", "qux", "quux")])
  
      def test_dont_remove_importfrom_simple_wrong_alias(self) -> None:
          """
          Should not remove import from since it is wrong alias
          """
  
          before = """
              import bar
              from baz import qux as quux
  
              def foo() -> None:
                  pass
          """
          after = """
              import bar
              from baz import qux as quux
  
              def foo() -> None:
                  pass
          """
  
          self.assertCodemod(before, after, [("baz", "qux", None)])
  
      def test_dont_remove_importfrom_alias_simple_wrong_alias(self) -> None:
          """
          Should not remove import from with wrong alias
          """
  
          before = """
              import bar
              from baz import qux
  
              def foo() -> None:
                  pass
          """
          after = """
              import bar
              from baz import qux
  
              def foo() -> None:
                  pass
          """
  
          self.assertCodemod(before, after, [("baz", "qux", "quux")])
  
      def test_remove_importfrom_relative(self) -> None:
          """
          Should remove import from which is relative
          """
  
          before = """
              import bar
              from .c import qux
  
              def foo() -> None:
                  pass
          """
          after = """
              import bar
  
              def foo() -> None:
                  pass
          """
  
          self.assertCodemod(
              before,
              after,
              [("a.b.c", "qux", None)],
              context_override=CodemodContext(full_module_name="a.b.foobar"),
          )
  
      def test_dont_remove_inuse_importfrom_relative(self) -> None:
          """
          Should not remove import from which is relative since it is in use.
          """
  
          before = """
              import bar
              from .c import qux
  
              def foo() -> None:
                  qux()
          """
          after = """
              import bar
              from .c import qux
  
              def foo() -> None:
                  qux()
          """
  
          self.assertCodemod(
              before,
              after,
              [("a.b.c", "qux", None)],
              context_override=CodemodContext(full_module_name="a.b.foobar"),
          )
  
      def test_dont_remove_wrong_importfrom_relative(self) -> None:
          """
          Should not remove import from which is relative since it is the wrong module.
          """
  
          before = """
              import bar
              from .c import qux
  
              def foo() -> None:
                  pass
          """
          after = """
              import bar
              from .c import qux
  
              def foo() -> None:
                  pass
          """
  
          self.assertCodemod(
              before,
              after,
              [("a.b.d", "qux", None)],
              context_override=CodemodContext(full_module_name="a.b.foobar"),
          )
  
      def test_remove_import_complex(self) -> None:
          """
          Should remove complex module as import
          """
  
          before = """
              import bar
              import baz, qux
              import a.b
              import c.d
              import x.y.z
              import e.f as g
              import h.i as j
  
              def foo() -> None:
                  c.d()
                  x.u
                  j()
          """
          after = """
              import bar
              import qux
              import c.d
              import x.y.z
              import h.i as j
  
              def foo() -> None:
                  c.d()
                  x.u
                  j()
          """
  
          self.assertCodemod(
              before,
              after,
              [
                  ("baz", None, None),
                  ("a.b", None, None),
                  ("c.d", None, None),
                  ("e.f", None, "g"),
                  ("h.i", None, "j"),
                  ("x.y.z", None, None),
              ],
          )
  
      def test_remove_fromimport_complex(self) -> None:
          """
          Should remove complex from import
          """
  
          before = """
              from bar import qux, quux
              from a.b import c
              from d.e import f
              from h.i import j as k
              from l.m import n as o
              from x import *
  
              def foo() -> None:
                  f()
                  k()
          """
          after = """
              from bar import qux
              from d.e import f
              from h.i import j as k
              from x import *
  
              def foo() -> None:
                  f()
                  k()
          """
  
          self.assertCodemod(
              before,
              after,
              [
                  ("bar", "quux", None),
                  ("a.b", "c", None),
                  ("d.e", "f", None),
                  ("h.i", "j", "k"),
                  ("l.m", "n", "o"),
              ],
          )
  
      def test_remove_import_multiple_assignments(self) -> None:
          """
          Should not remove import with multiple assignments
          """
  
          before = """
              from foo import bar
              from qux import bar
  
              def foo() -> None:
                  bar()
          """
          after = """
              from foo import bar
              from qux import bar
  
              def foo() -> None:
                  bar()
          """
  
          self.assertCodemod(before, after, [("foo", "bar", None)])
  
      def test_remove_multiple_imports(self) -> None:
          """
          Multiple imports
          """
          before = """
              try:
                  import a
              except Exception:
                  import a
  
              a.hello()
          """
          after = """
              try:
                  import a
              except Exception:
                  import a
  
              a.hello()
          """
          self.assertCodemod(before, after, [("a", None, None)])
  
          before = """
              try:
                  import a
              except Exception:
                  import a
          """
          after = """
              try:
                  pass
              except Exception:
                  pass
          """
          self.assertCodemod(before, after, [("a", None, None)])
  
      @data_provider(
          (
              # Simple removal, no other uses.
              (
                  """
                      from foo import bar
                      from qux import baz
  
                      def fun() -> None:
                          bar()
                          baz()
                  """,
                  """
                      from qux import baz
  
                      def fun() -> None:
                          baz()
                  """,
              ),
              # Remove a node, other uses, don't remove import.
              (
                  """
                      from foo import bar
                      from qux import baz
  
                      def fun() -> None:
                          bar()
                          baz()
  
                      def foobar() -> None:
                          a = bar
                          a()
                  """,
                  """
                      from foo import bar
                      from qux import baz
  
                      def fun() -> None:
                          baz()
  
                      def foobar() -> None:
                          a = bar
                          a()
                  """,
              ),
              # Remove an alias.
              (
                  """
                      from foo import bar as other
                      from qux import baz
  
                      def fun() -> None:
                          other()
                          baz()
                  """,
                  """
                      from qux import baz
  
                      def fun() -> None:
                          baz()
                  """,
              ),
              # Simple removal, no other uses.
              (
                  """
                      import foo
                      from qux import baz
  
                      def fun() -> None:
                          foo.bar()
                          baz()
                  """,
                  """
                      from qux import baz
  
                      def fun() -> None:
                          baz()
                  """,
              ),
              # Remove a node, other uses, don't remove import.
              (
                  """
                      import foo
                      from qux import baz
  
                      def fun() -> None:
                          foo.bar()
                          baz()
  
                      def foobar() -> None:
                          a = foo.bar
                          a()
                  """,
                  """
                      import foo
                      from qux import baz
  
                      def fun() -> None:
                          baz()
  
                      def foobar() -> None:
                          a = foo.bar
                          a()
                  """,
              ),
              # Remove an alias.
              (
                  """
                      import foo as other
                      from qux import baz
  
                      def fun() -> None:
                          other.bar()
                          baz()
                  """,
                  """
                      from qux import baz
  
                      def fun() -> None:
                          baz()
                  """,
              ),
          )
      )
      def test_remove_import_by_node_simple(self, before: str, after: str) -> None:
          """
          Given a node that's directly referenced in an import,
          make sure that the import is removed when the node
          is also removed.
          """
  
          class RemoveBarTransformer(VisitorBasedCodemodCommand):
  
              METADATA_DEPENDENCIES = (QualifiedNameProvider, ScopeProvider)
  
              @m.leave(
                  m.SimpleStatementLine(
                      body=[
                          m.Expr(
                              m.Call(
                                  metadata=m.MatchMetadata(
                                      QualifiedNameProvider,
                                      {
                                          QualifiedName(
                                              source=QualifiedNameSource.IMPORT,
                                              name="foo.bar",
                                          )
                                      },
                                  )
                              )
                          )
                      ]
                  )
              )
              def _leave_foo_bar(
                  self,
                  original_node: cst.SimpleStatementLine,
                  updated_node: cst.SimpleStatementLine,
              ) -> cst.RemovalSentinel:
                  RemoveImportsVisitor.remove_unused_import_by_node(
                      self.context, original_node
                  )
                  return cst.RemoveFromParent()
  
          module = cst.parse_module(self.make_fixture_data(before))
          self.assertCodeEqual(
              after, RemoveBarTransformer(CodemodContext()).transform_module(module).code
          )
  
      def test_remove_import_from_node(self) -> None:
          """
          Make sure that if an import node itself is requested for
          removal, we still do the right thing and only remove it
          if it is unused.
          """
  
          before = """
              from foo import bar
              from qux import baz
              from foo import qux as other
              from qux import foobar as other2
  
              def fun() -> None:
                  baz()
                  other2()
          """
          after = """
              from qux import baz
              from qux import foobar as other2
  
              def fun() -> None:
                  baz()
                  other2()
          """
  
          class RemoveImportTransformer(VisitorBasedCodemodCommand):
  
              METADATA_DEPENDENCIES = (QualifiedNameProvider, ScopeProvider)
  
              def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
                  RemoveImportsVisitor.remove_unused_import_by_node(self.context, node)
  
          module = cst.parse_module(self.make_fixture_data(before))
          self.assertCodeEqual(
              after,
              RemoveImportTransformer(CodemodContext()).transform_module(module).code,
          )
  
      def test_remove_import_node(self) -> None:
          """
          Make sure that if an import node itself is requested for
          removal, we still do the right thing and only remove it
          if it is unused.
          """
  
          before = """
              import foo
              import qux
              import bar as other
              import foobar as other2
  
              def fun() -> None:
                  qux.baz()
                  other2.baz()
          """
          after = """
              import qux
              import foobar as other2
  
              def fun() -> None:
                  qux.baz()
                  other2.baz()
          """
  
          class RemoveImportTransformer(VisitorBasedCodemodCommand):
  
              METADATA_DEPENDENCIES = (QualifiedNameProvider, ScopeProvider)
  
              def visit_Import(self, node: cst.Import) -> None:
                  RemoveImportsVisitor.remove_unused_import_by_node(self.context, node)
  
          module = cst.parse_module(self.make_fixture_data(before))
          self.assertCodeEqual(
              after,
              RemoveImportTransformer(CodemodContext()).transform_module(module).code,
          )
  
      def test_remove_import_with_all(self) -> None:
          """
          Make sure that if an import node itself is requested for
          removal, we don't remove it if it shows up in an __all__
          node.
          """
  
          before = """
              from foo import bar
              from qux import baz
  
              __all__ = ["baz"]
          """
          after = """
              from qux import baz
  
              __all__ = ["baz"]
          """
  
          class RemoveImportTransformer(VisitorBasedCodemodCommand):
  
              METADATA_DEPENDENCIES = (QualifiedNameProvider, ScopeProvider)
  
              def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
                  RemoveImportsVisitor.remove_unused_import_by_node(self.context, node)
  
          module = cst.parse_module(self.make_fixture_data(before))
          self.assertCodeEqual(
              after,
              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.
          """
          before = """
              from m import (a, b,)
              import x, y
          """
          after = """
              from m import (b,)
              import x
          """
          self.assertCodemod(before, after, [("m", "a", None), ("y", None, None)])
  