Tawsif fix asyncify promise return type

src/codegen/sdk/typescript/function.py

@@ -34,7 +34,7 @@
 
 
 @ts_apidoc
 class TSFunction(Function["TSFunction", TSDecorator, "TSCodeBlock", TSParameter, TSType], TSHasBlock, TSSymbol):
     """Representation of a Function in JavaScript/TypeScript"""
 
     @noapidoc
@@ -44,7 +44,7 @@
 
         self.return_type = self.child_by_field_name("return_type", placeholder=TSReturnTypePlaceholder)
         if parameters_node := self.ts_node.child_by_field_name("parameters"):
             self._parameters = Collection(parameters_node, self.file_node_id, self.G, self)
             params = [x for x in parameters_node.children if x.type in ("required_parameter", "optional_parameter")]
             symbols = None
             # Deconstructed object parameters
@@ -55,18 +55,18 @@
                     type_annotation = self._parse_type(type_node)
                 if pattern and pattern.type == "object_pattern":
                     params = [x for x in pattern.children if x.type in ("shorthand_property_identifier_pattern", "object_assignment_pattern", "pair_pattern")]
                     symbols = [TSParameter(x, i, self._parameters, type_annotation) for (i, x) in enumerate(params)]
             # Default case - regular parameters
             if symbols is None:
                 symbols = [TSParameter(x, i, self._parameters) for (i, x) in enumerate(params)]
             self._parameters._init_children(symbols)
         elif parameters_node := self.ts_node.child_by_field_name("parameter"):
             self._parameters = Collection(parameters_node, self.file_node_id, self.G, self)
             self._parameters._init_children([TSParameter(parameters_node, 0, self._parameters)])
         else:
             logger.warning(f"Couldn't find parameters for {self!r}")
             self._parameters = []
         self.type_parameters = self.child_by_field_name("type_parameters")
 
     @property
     @reader
@@ -97,7 +97,7 @@
             for param in self.parameters:
                 assignment_patterns = find_all_descendants(param.ts_node, {"object_pattern", "object_assignment_pattern", "assignment_pattern"})
                 if assignment_patterns:
                     dest.add_all_identifier_usages_for_child_node(UsageKind.GENERIC, assignment_patterns[0])
         if self.type_parameters:
             self.type_parameters._compute_dependencies(UsageKind.GENERIC, dest)
         # =====[ Return type ]=====
@@ -106,7 +106,7 @@
             self.return_type._compute_dependencies(UsageKind.RETURN_TYPE, dest)
 
         # =====[ Code Block ]=====
         self.code_block._compute_dependencies(usage_type, dest)
 
     @classmethod
     @noapidoc
@@ -297,7 +297,7 @@
         if self.is_async:
             return
         self.add_keyword("async")
-        if self.return_type:
+        if self.return_type and self.return_type.name != "Promise":
             self.return_type.insert_before("Promise<", newline=False)
             self.return_type.insert_after(">", newline=False)
 

tests/unit/codegen/sdk/typescript/function/test_function_async.py

@@ -1,5 +1,6 @@
 from codegen.sdk.codebase.factory.get_session import get_codebase_session
 from codegen.sdk.enums import ProgrammingLanguage
+from codegen.sdk.typescript.placeholder.placeholder_return_type import TSReturnTypePlaceholder
 
 
 def test_function_is_async_basic(tmpdir):
@@ -233,3 +234,149 @@ class MathOperations {
 }
     """
     )
+
+
+def test_asyncify_wraps_non_promise_return_type(tmpdir) -> None:
+    # ========= = [ BEFORE ] ==========
+    # language=typescript
+    BEFORE_CONTENT = """
+function getData(): string {
+    return "hello";
+}
+"""
+    # ========== [ AFTER ] ==========
+    # language=typescript
+    EXPECTED_CONTENT = """
+async function getData(): Promise<string> {
+    return "hello";
+}
+"""
+
+    with get_codebase_session(
+        tmpdir=tmpdir,
+        programming_language=ProgrammingLanguage.TYPESCRIPT,
+        files={"test.ts": BEFORE_CONTENT},
+    ) as codebase:
+        file = codebase.get_file("test.ts")
+        func = file.get_function("getData")
+
+        # Initial state should be non-async
+        assert not func.is_async
+        assert func.return_type.source == "string"
+
+        # After asyncify, should be async and return type wrapped in Promise
+        func.asyncify()
+        codebase.commit()
+
+        # Check file content directly instead of func.is_async
+        assert file.content.strip() == EXPECTED_CONTENT.strip()
+
+
+def test_asyncify_already_promise_return_type(tmpdir) -> None:
+    # ========== [ BEFORE ] ==========
+    # language=typescript
+    BEFORE_CONTENT = """
+    function getData(): Promise<string> {
+        return Promise.resolve("hello");
+    }
+    """
+
+    # ========== [ AFTER ] ==========
+    # language=typescript
+    EXPECTED_CONTENT = """
+    async function getData(): Promise<string> {
+        return Promise.resolve("hello");
+    }
+    """
+
+    with get_codebase_session(
+        tmpdir=tmpdir,
+        programming_language=ProgrammingLanguage.TYPESCRIPT,
+        files={"test.ts": BEFORE_CONTENT},
+    ) as codebase:
+        file = codebase.get_file("test.ts")
+        func = file.get_function("getData")
+
+        # Initial state should be non-async but already have Promise return type
+        assert not func.is_async
+        assert func.return_type.source == "Promise<string>"
+
+        # After asyncify, should be async but return type should remain unchanged
+        func.asyncify()
+        codebase.commit()
+
+        # Check file content directly instead of func.is_async
+        print(file.content)
+        assert file.content.strip() == EXPECTED_CONTENT.strip()
+
+
+def test_asyncify_void_return_type(tmpdir) -> None:
+    # ========== [ BEFORE ] ==========
+    # language=typescript
+    BEFORE_CONTENT = """
+    function processData(): void {
+        console.log("processing");
+    }
+    """
+
+    # ========== [ AFTER ] ==========
+    # language=typescript
+    EXPECTED_CONTENT = """
+    async function processData(): Promise<void> {
+        console.log("processing");
+    }
+    """
+
+    with get_codebase_session(
+        tmpdir=tmpdir,
+        programming_language=ProgrammingLanguage.TYPESCRIPT,
+        files={"test.ts": BEFORE_CONTENT},
+    ) as codebase:
+        file = codebase.get_file("test.ts")
+        func = file.get_function("processData")
+
+        # Initial state should be non-async with void return type
+        assert not func.is_async
+        assert func.return_type.source == "void"
+
+        # After asyncify, should be async and return Promise<void>
+        func.asyncify()
+        codebase.commit()
+        # Check file content directly instead of func.is_async
+        assert file.content.strip() == EXPECTED_CONTENT.strip()
+
+
+def test_asyncify_no_return_type(tmpdir) -> None:
+    # ========== [ BEFORE ] ==========
+    # language=typescript
+    BEFORE_CONTENT = """
+    function processData() {
+        console.log("processing");
+    }
+    """
+
+    # ========== [ AFTER ] ==========
+    # language=typescript
+    EXPECTED_CONTENT = """
+    async function processData() {
+        console.log("processing");
+    }
+    """
+
+    with get_codebase_session(
+        tmpdir=tmpdir,
+        programming_language=ProgrammingLanguage.TYPESCRIPT,
+        files={"test.ts": BEFORE_CONTENT},
+    ) as codebase:
+        file = codebase.get_file("test.ts")
+        func = file.get_function("processData")
+
+        # Initial state should be non-async with no return type
+        assert not func.is_async
+        assert isinstance(func.return_type, TSReturnTypePlaceholder)
+
+        # After asyncify, should be async but no return type added
+        func.asyncify()
+        codebase.commit()
+        # Check file content directly instead of func.is_async
+        assert file.content.strip() == EXPECTED_CONTENT.strip()