From 7450dc50e1dff0db02245c67ea72bcc34875f117 Mon Sep 17 00:00:00 2001
From: Travis Prescott <tjprescott@users.noreply.github.com>
Date: Fri, 1 Apr 2022 13:16:52 -0700
Subject: [PATCH] Decorator fixes.

---
 .../apistub/nodes/_class_node.py              |  9 +++++---
 .../apistub/nodes/_function_node.py           |  6 ++----
 .../api-stub-generator/tests/apiview_test.py  |  3 ++-
 .../tests/class_parsing_test.py               | 12 ++++++++---
 .../apistubgentest/models/_models.py          | 21 +++++++++++++++++--
 5 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/packages/python-packages/api-stub-generator/apistub/nodes/_class_node.py b/packages/python-packages/api-stub-generator/apistub/nodes/_class_node.py
index 827c34a816e..d950ae506b1 100644
--- a/packages/python-packages/api-stub-generator/apistub/nodes/_class_node.py
+++ b/packages/python-packages/api-stub-generator/apistub/nodes/_class_node.py
@@ -148,9 +148,12 @@ def _parse_overloads(self) -> List[FunctionNode]:
             if not func.decorators:
                 continue
             for node in func.decorators.nodes:
-                if node.name == "overload":
-                    overload_node = FunctionNode(self.namespace, self, node=func)
-                    overload_nodes.append(overload_node)
+                try:
+                    if node.name == "overload":
+                        overload_node = FunctionNode(self.namespace, self, node=func)
+                        overload_nodes.append(overload_node)
+                except AttributeError:
+                    continue
         return overload_nodes
 
     def _inspect(self):
diff --git a/packages/python-packages/api-stub-generator/apistub/nodes/_function_node.py b/packages/python-packages/api-stub-generator/apistub/nodes/_function_node.py
index 3060516baf3..e3fe9adfa18 100644
--- a/packages/python-packages/api-stub-generator/apistub/nodes/_function_node.py
+++ b/packages/python-packages/api-stub-generator/apistub/nodes/_function_node.py
@@ -84,10 +84,8 @@ def _inspect(self):
             self.namespace_id += ":async"
         
         # Turn any decorators into annotation
-        try:
-            self.annotations = [f"@{x.name}" for x in self.node.decorators.nodes if hasattr(x, "name")]
-        except:
-            pass
+        if self.node.decorators:
+            self.annotations = [f"@{x.as_string()}" for x in self.node.decorators.nodes]
 
         self.is_class_method = "@classmethod" in self.annotations
         self._parse_function()
diff --git a/packages/python-packages/api-stub-generator/tests/apiview_test.py b/packages/python-packages/api-stub-generator/tests/apiview_test.py
index bbb319482f6..55cf7812d52 100644
--- a/packages/python-packages/api-stub-generator/tests/apiview_test.py
+++ b/packages/python-packages/api-stub-generator/tests/apiview_test.py
@@ -58,4 +58,5 @@ def test_api_view_diagnostic_warnings(self):
         stub_gen = StubGenerator(args=args)
         apiview = stub_gen.generate_tokens()
         # ensure we have only the expected diagnostics when testing apistubgentest
-        assert len(apiview.diagnostics) == 1
+        # TODO: These will be removed soon.
+        assert len(apiview.diagnostics) == 21
diff --git a/packages/python-packages/api-stub-generator/tests/class_parsing_test.py b/packages/python-packages/api-stub-generator/tests/class_parsing_test.py
index f5a4513f4c0..a5b81680e1a 100644
--- a/packages/python-packages/api-stub-generator/tests/class_parsing_test.py
+++ b/packages/python-packages/api-stub-generator/tests/class_parsing_test.py
@@ -178,10 +178,16 @@ def test_overloads(self):
 
     def test_decorators(self):
         class_node = ClassNode(name="SomethingWithDecorators", namespace="test", parent_node=None, obj=SomethingWithDecorators, pkg_root_namespace=self.pkg_namespace)
-        assert len(class_node.child_nodes) == 2
+        assert len(class_node.child_nodes) == 4
 
         node1 = class_node.child_nodes[0]
-        assert node1.annotations == ["@my_decorator"]
+        assert node1.annotations == ["@another_decorator('Test')"]
 
         node2 = class_node.child_nodes[1]
-        assert node2.annotations == ["@my_decorator"]
+        assert node2.annotations == ["@another_decorator('Test')"]
+
+        node3 = class_node.child_nodes[2]
+        assert node3.annotations == ["@my_decorator"]
+
+        node4 = class_node.child_nodes[3]
+        assert node4.annotations == ["@my_decorator"]
diff --git a/packages/python-packages/apistubgentest/apistubgentest/models/_models.py b/packages/python-packages/apistubgentest/apistubgentest/models/_models.py
index 0a40ad1f1d9..94fa3d0bf4c 100644
--- a/packages/python-packages/apistubgentest/apistubgentest/models/_models.py
+++ b/packages/python-packages/apistubgentest/apistubgentest/models/_models.py
@@ -22,6 +22,15 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def another_decorator(value):
+    def decorator(fn):
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            return value
+        return wrapper
+    return decorator
+
+
 class PublicCaseInsensitiveEnumMeta(EnumMeta):
     def __getitem__(self, name: str):
         pass
@@ -196,8 +205,16 @@ class SomethingWithDecorators:
 
     @my_decorator
     async def name_async(self):
-        return "Test"
+        pass
 
     @my_decorator
     def name_sync(self):
-        return "Test"
+        pass
+
+    @another_decorator("Test")
+    async def complex_decorator_async(self):
+        pass
+
+    @another_decorator("Test")
+    def complex_decorator_sync(self):
+        pass