From 3009ec9e36e5bdca7f284bcac3d4fcc810b4a922 Mon Sep 17 00:00:00 2001
From: Steven Troxler <steven.troxler@gmail.com>
Date: Fri, 6 Aug 2021 08:55:47 -0700
Subject: [PATCH] Add tests verifying TypeAnnotationsVisitor current behaviors
 (#502)

* Add tests to verify that LibCST handles string annotations.

This is an important property for certain use cases, so it
makes sense to verify it in tests so that we can safely
depend on it.

At present, the reason we want to be able to rely on this is:
- at the moment, imports added by infer can make pysa traces
  hard to understand, because the line numbers are off
- if we add the ability to use fully-qualified string annotations
  for the stubs from infer, then we can do so without adding
  any import lines and pyre will understand the types.

* ApplyTypeAnnotations: add unit test of how import statments are merged

Add a unit test illustrating how the codemod handles various cases
of import statments in the stub file.

Explicitly call out each of the unsupported patterns:
- bare imports (we probably should support this)
- relative imports (we probably should support this)
 star imports (we probably don't want to support this)

* Add .python-version to .gitignore

This will be helpful for anyone using pyenv (I accidentally committed
my python version file in a draft branch).
---
 .gitignore                                    |   1 +
 .../tests/test_apply_type_annotations.py      | 111 +++++++++++++++---
 2 files changed, 96 insertions(+), 16 deletions(-)

diff --git a/.gitignore b/.gitignore
index 85fb5573a..2c52df24f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,3 +13,4 @@ build/
 .coverage
 .hypothesis/
 .pyre_configuration
+.python-version
diff --git a/libcst/codemod/visitors/tests/test_apply_type_annotations.py b/libcst/codemod/visitors/tests/test_apply_type_annotations.py
index d41ae8f58..16ccda993 100644
--- a/libcst/codemod/visitors/tests/test_apply_type_annotations.py
+++ b/libcst/codemod/visitors/tests/test_apply_type_annotations.py
@@ -18,6 +18,76 @@
 class TestApplyAnnotationsVisitor(CodemodTest):
     TRANSFORM: Type[Codemod] = ApplyTypeAnnotationsVisitor
 
+    @data_provider(
+        (
+            (
+                """
+                from __future__ import annotations
+                from foo import Foo
+                from baz import Baz
+                """,
+                """
+                from foo import Bar
+                import bar
+                """,
+                """
+                from __future__ import annotations
+                from foo import Foo, Bar
+                import bar
+                from baz import Baz
+                """,
+            ),
+            (
+                # Missing feature: ignore aliased imports
+                """
+                from Foo import foo as bar
+                """,
+                """
+                from Foo import bar
+                """,
+                """
+                from Foo import bar
+                """,
+            ),
+            (
+                # Missing feature: ignore bare imports
+                """
+                import foo
+                """,
+                """
+                """,
+                """
+                """,
+            ),
+            (
+                # Missing feature: ignore relative imports
+                """
+                from .. import foo
+                """,
+                """
+                """,
+                """
+                """,
+            ),
+            (
+                # Missing feature: ignore star imports
+                """
+                from foo import *
+                """,
+                """
+                """,
+                """
+                """,
+            ),
+        )
+    )
+    def test_merge_module_imports(self, stub: str, before: str, after: str) -> None:
+        context = CodemodContext()
+        ApplyTypeAnnotationsVisitor.store_stub_in_context(
+            context, parse_module(textwrap.dedent(stub.rstrip()))
+        )
+        self.assertCodemod(before, after, context_override=context)
+
     @data_provider(
         (
             (
@@ -608,22 +678,6 @@ def foo() -> None:
                     pass
                 """,
             ),
-            # Sanity check that we don't fail when the stub has relative imports.
-            # We don't do anything with those imports, though.
-            (
-                """
-                from .. import hello
-                def foo() -> typing.Sequence[int]: ...
-                """,
-                """
-                def foo():
-                  return []
-                """,
-                """
-                def foo() -> typing.Sequence[int]:
-                  return []
-                """,
-            ),
             (
                 """
                 from typing import Dict
@@ -669,6 +723,31 @@ class A:
                     def foo(self, atticus, b: Optional[int] = None, c: bool = False): ...
                 """,
             ),
+            # Make sure we handle string annotations well
+            (
+                """
+                def f(x: "typing.Union[int, str]") -> "typing.Union[int, str]": ...
+
+                class A:
+                    def f(self: "A") -> "A": ...
+                """,
+                """
+                def f(x):
+                    return x
+
+                class A:
+                    def f(self):
+                        return self
+                """,
+                """
+                def f(x: "typing.Union[int, str]") -> "typing.Union[int, str]":
+                    return x
+
+                class A:
+                    def f(self: "A") -> "A":
+                        return self
+                """,
+            ),
         )
     )
     def test_annotate_functions(self, stub: str, before: str, after: str) -> None: