From d9e04544c3e1e280414de72a9dc0ae27977fb9b9 Mon Sep 17 00:00:00 2001
From: Tianning Li <ltianningli@gmail.com>
Date: Mon, 12 Feb 2024 16:05:53 -0800
Subject: [PATCH] feat: add error response decorator

---
 lib/plugin/utils/ast-utils.ts                 |  29 +++++
 .../visitors/controller-class.visitor.ts      | 100 +++++++++++++++++-
 2 files changed, 124 insertions(+), 5 deletions(-)

diff --git a/lib/plugin/utils/ast-utils.ts b/lib/plugin/utils/ast-utils.ts
index 5eeea754b..9387df043 100644
--- a/lib/plugin/utils/ast-utils.ts
+++ b/lib/plugin/utils/ast-utils.ts
@@ -228,6 +228,35 @@ export function getTsDocTagsOfNode(node: Node, typeChecker: TypeChecker) {
   return tagResults;
 }
 
+export function getTsDocReturnsOrErrorOfNode(node: Node) {
+  const tsdocParser: TSDocParser = new TSDocParser();
+  const parserContext: ParserContext = tsdocParser.parseString(
+    node.getFullText()
+  );
+  const docComment: DocComment = parserContext.docComment;
+
+  const tagResults = [];
+  const introspectTsDocTags = (docComment: DocComment) => {
+    const blocks = docComment.customBlocks.filter((block) =>
+      ['@throws', '@returns'].includes(block.blockTag.tagName)
+    );
+
+    blocks.forEach((block) => {
+      try {
+        const docValue = renderDocNode(block.content).split('\n')[0].trim();
+        const regex = /{(\d+)} (.*)/;
+        const match = docValue.match(regex);
+        tagResults.push({
+          status: match[1],
+          description: `"${match[2]}"`
+        });
+      } catch (err) {}
+    });
+  };
+  introspectTsDocTags(docComment);
+  return tagResults;
+}
+
 export function getDecoratorArguments(decorator: Decorator) {
   const callExpression = decorator.expression;
   return (callExpression && (callExpression as CallExpression).arguments) || [];
diff --git a/lib/plugin/visitors/controller-class.visitor.ts b/lib/plugin/visitors/controller-class.visitor.ts
index 68162c6e5..794707e7b 100644
--- a/lib/plugin/visitors/controller-class.visitor.ts
+++ b/lib/plugin/visitors/controller-class.visitor.ts
@@ -9,6 +9,7 @@ import {
   getDecoratorArguments,
   getDecoratorName,
   getMainCommentOfNode,
+  getTsDocReturnsOrErrorOfNode,
   getTsDocTagsOfNode
 } from '../utils/ast-utils';
 import {
@@ -139,14 +140,34 @@ export class ControllerClassVisitor extends AbstractFileVisitor {
       typeChecker,
       metadata
     );
+
+    const apiResponseDecoratorsArray = this.createApiResponseDecorator(
+      factory,
+      compilerNode,
+      decorators,
+      options,
+      sourceFile,
+      typeChecker,
+      metadata
+    );
+
     const removeExistingApiOperationDecorator =
       apiOperationDecoratorsArray.length > 0;
 
-    const existingDecorators = removeExistingApiOperationDecorator
-      ? decorators.filter(
-          (item) => getDecoratorName(item) !== ApiOperation.name
-        )
-      : decorators;
+    const removeExistingApiResponseDecorator =
+      apiResponseDecoratorsArray.length > 0;
+
+    let existingDecorators = decorators;
+    if (
+      removeExistingApiOperationDecorator ||
+      removeExistingApiResponseDecorator
+    ) {
+      existingDecorators = decorators.filter(
+        (item) =>
+          getDecoratorName(item) !== ApiOperation.name &&
+          getDecoratorName(item) !== ApiResponse.name
+      );
+    }
 
     const modifiers = ts.getModifiers(compilerNode) ?? [];
     const objectLiteralExpr = this.createDecoratorObjectLiteralExpr(
@@ -160,6 +181,7 @@ export class ControllerClassVisitor extends AbstractFileVisitor {
     );
     const updatedDecorators = [
       ...apiOperationDecoratorsArray,
+      ...apiResponseDecoratorsArray,
       ...existingDecorators,
       factory.createDecorator(
         factory.createCallExpression(
@@ -302,6 +324,74 @@ export class ControllerClassVisitor extends AbstractFileVisitor {
     }
   }
 
+  createApiResponseDecorator(
+    factory: ts.NodeFactory,
+    node: ts.MethodDeclaration,
+    decorators: readonly ts.Decorator[],
+    options: PluginOptions,
+    sourceFile: ts.SourceFile,
+    typeChecker: ts.TypeChecker,
+    metadata: ClassMetadata
+  ) {
+    if (!options.introspectComments) {
+      return [];
+    }
+    const apiResponseDecorator = getDecoratorOrUndefinedByNames(
+      [ApiResponse.name],
+      decorators,
+      factory
+    );
+    let apiResponseExistingProps:
+      | ts.NodeArray<ts.PropertyAssignment>
+      | undefined = undefined;
+
+    if (apiResponseDecorator && !options.readonly) {
+      const apiResponseExpr = head(getDecoratorArguments(apiResponseDecorator));
+      if (apiResponseExpr) {
+        apiResponseExistingProps =
+          apiResponseExpr.properties as ts.NodeArray<ts.PropertyAssignment>;
+      }
+    }
+
+    const tags = getTsDocReturnsOrErrorOfNode(node);
+    if (!tags.length) {
+      return [];
+    }
+
+    return tags.map((tag) => {
+      const properties = [
+        ...(apiResponseExistingProps ?? factory.createNodeArray())
+      ];
+      properties.push(
+        factory.createPropertyAssignment(
+          'status',
+          factory.createNumericLiteral(tag.status)
+        )
+      );
+      properties.push(
+        factory.createPropertyAssignment(
+          'description',
+          factory.createNumericLiteral(tag.description)
+        )
+      );
+      const objectLiteralExpr = factory.createObjectLiteralExpression(
+        compact(properties)
+      );
+      const methodKey = node.name.getText();
+      metadata[methodKey] = objectLiteralExpr;
+
+      const apiResponseDecoratorArguments: ts.NodeArray<ts.Expression> =
+        factory.createNodeArray([objectLiteralExpr]);
+      return factory.createDecorator(
+        factory.createCallExpression(
+          factory.createIdentifier(`${OPENAPI_NAMESPACE}.${ApiResponse.name}`),
+          undefined,
+          apiResponseDecoratorArguments
+        )
+      );
+    });
+  }
+
   createDecoratorObjectLiteralExpr(
     factory: ts.NodeFactory,
     node: ts.MethodDeclaration,