stubber: Also merge TypeAliases

Signed-off-by: Jos Verlinde <Jos.Verlinde@microsoft.com>

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):

tests/codemods/codemod_test_cases/typealias_add/before.py

@@ -0,0 +1,10 @@
+#fmt: off
+"""
+add typealias
+"""
+
+
+def const(): ...
+
+
+def bar(): ...

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::
+    """
+    ...

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::
+    """
+    ...