From eb6db3ae61e13d8ad527796e713695dac9aa052a Mon Sep 17 00:00:00 2001
From: Jos Verlinde <Jos.Verlinde@microsoft.com>
Date: Sun, 27 Oct 2024 21:28:37 +0100
Subject: [PATCH] stubber: Also merge TypeAliases

Signed-off-by: Jos Verlinde <Jos.Verlinde@microsoft.com>
---
 src/stubber/codemod/visitors/typevars.py      | 16 +++++++++-----
 .../typealias_add/before.py                   | 10 +++++++++
 .../typealias_add/doc_stub.py                 | 21 ++++++++++++++++++
 .../typealias_add/expected.py                 | 22 +++++++++++++++++++
 4 files changed, 64 insertions(+), 5 deletions(-)
 create mode 100644 tests/codemods/codemod_test_cases/typealias_add/before.py
 create mode 100644 tests/codemods/codemod_test_cases/typealias_add/doc_stub.py
 create mode 100644 tests/codemods/codemod_test_cases/typealias_add/expected.py

diff --git a/src/stubber/codemod/visitors/typevars.py b/src/stubber/codemod/visitors/typevars.py
index dfe432b6..b0374a6d 100644
--- a/src/stubber/codemod/visitors/typevars.py
+++ b/src/stubber/codemod/visitors/typevars.py
@@ -24,16 +24,22 @@ def __init__(self, context: CodemodContext) -> None:
         self.all_typevars: List[Union[libcst.Assign, libcst.AnnAssign]] = []
 
     def visit_Assign(self, node: libcst.Assign) -> None:
-
+        """
+        Find all TypeVar assignments in the module.
+        format: T = TypeVar("T", int, float, str, bytes, Tuple)
+        """
         # is this a TypeVar assignment?
         # needs to be more robust
-
         if isinstance(node.value, libcst.Call) and node.value.func.value == "TypeVar":
             self.all_typevars.append(node)
 
-    # def visit_AnnAssign(self, node: libcst.AnnAssign) -> None:
-    #     # Track this import statement for later analysis.
-    #     self.all_typevars.append(node)
+    def visit_AnnAssign(self, node: libcst.AnnAssign) -> None:
+        """ "
+        Find all TypeAlias assignments in the module.
+        format: T: TypeAlias = str
+        """
+        if isinstance(node.annotation.annotation, libcst.Name) and node.annotation.annotation.value == "TypeAlias":
+            self.all_typevars.append(node)
 
 
 class AddTypeVarsVisitor(ContextAwareTransformer):
diff --git a/tests/codemods/codemod_test_cases/typealias_add/before.py b/tests/codemods/codemod_test_cases/typealias_add/before.py
new file mode 100644
index 00000000..9ef19cb8
--- /dev/null
+++ b/tests/codemods/codemod_test_cases/typealias_add/before.py
@@ -0,0 +1,10 @@
+#fmt: off
+"""
+add typealias
+"""
+
+
+def const(): ...
+
+
+def bar(): ...
diff --git a/tests/codemods/codemod_test_cases/typealias_add/doc_stub.py b/tests/codemods/codemod_test_cases/typealias_add/doc_stub.py
new file mode 100644
index 00000000..35583220
--- /dev/null
+++ b/tests/codemods/codemod_test_cases/typealias_add/doc_stub.py
@@ -0,0 +1,21 @@
+# fmt: off
+"""
+add typealias
+"""
+
+
+from typing import TypeAlias
+
+UUID: TypeAlias = str
+_Flag: TypeAlias = int
+_Descriptor: TypeAlias = tuple["UUID", _Flag]
+_Characteristic: TypeAlias = (tuple["UUID", _Flag] | tuple["UUID", _Flag, tuple[_Descriptor, ...]])
+_Service: TypeAlias = tuple["UUID", tuple[_Characteristic, ...]]
+x = 2
+
+def bar(f:_Flag , d:_Descriptor   ) -> _Service:
+    """
+    Used to declare that the expression is a constant so that the compiler can
+    optimise it.  The use of this function should be as follows::
+    """
+    ...
diff --git a/tests/codemods/codemod_test_cases/typealias_add/expected.py b/tests/codemods/codemod_test_cases/typealias_add/expected.py
new file mode 100644
index 00000000..7f61917b
--- /dev/null
+++ b/tests/codemods/codemod_test_cases/typealias_add/expected.py
@@ -0,0 +1,22 @@
+# fmt: off
+"""
+add typealias
+"""
+from typing import TypeAlias
+
+UUID: TypeAlias = str
+_Flag: TypeAlias = int
+_Descriptor: TypeAlias = tuple["UUID", _Flag]
+_Characteristic: TypeAlias = (tuple["UUID", _Flag] | tuple["UUID", _Flag, tuple[_Descriptor, ...]])
+_Service: TypeAlias = tuple["UUID", tuple[_Characteristic, ...]]
+
+
+def const(): ...
+
+
+def bar(f:_Flag , d:_Descriptor   ) -> _Service:
+    """
+    Used to declare that the expression is a constant so that the compiler can
+    optimise it.  The use of this function should be as follows::
+    """
+    ...