diff --git a/src/codegen/sdk/python/symbol.py b/src/codegen/sdk/python/symbol.py index d173b6cf9..3fd84631a 100644 --- a/src/codegen/sdk/python/symbol.py +++ b/src/codegen/sdk/python/symbol.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Self, Unpack from codegen.sdk.core.autocommit import reader, writer from codegen.sdk.core.symbol import Symbol @@ -13,6 +13,8 @@ from tree_sitter import Node as TSNode from codegen.sdk.codebase.codebase_graph import CodebaseGraph + from codegen.sdk.codebase.flagging.code_flag import CodeFlag + from codegen.sdk.codebase.flagging.enums import FlagKwargs from codegen.sdk.core.interfaces.has_block import HasBlock from codegen.sdk.core.node_id_factory import NodeId from codegen.sdk.python.detached_symbols.code_block import PyCodeBlock @@ -197,3 +199,26 @@ def set_inline_comment(self, comment: str, auto_format: bool = True, clean_forma if auto_format: comment = " " + PyComment.generate_comment(comment, PyCommentType.SINGLE_LINE) self.insert_after(comment, fix_indentation=False, newline=False) + + @writer + def flag(self, **kwargs: Unpack[FlagKwargs]) -> CodeFlag[Self]: + """Flags a Python symbol by adding a flag comment and returning a CodeFlag. + + This implementation first creates the CodeFlag through the standard flagging system, + then adds a Python-specific comment to visually mark the flagged code. + + Args: + **kwargs: Flag keyword arguments including optional 'message' + + Returns: + CodeFlag[Self]: The code flag object for tracking purposes + """ + # First create the standard CodeFlag through the base implementation + code_flag = super().flag(**kwargs) + + # Add a Python comment to visually mark the flag + message = kwargs.get("message", "") + if message: + self.set_inline_comment(f"🚩 {message}") + + return code_flag diff --git a/src/codegen/sdk/typescript/symbol.py b/src/codegen/sdk/typescript/symbol.py index 7ff4f0a47..bb41f3a52 100644 --- a/src/codegen/sdk/typescript/symbol.py +++ b/src/codegen/sdk/typescript/symbol.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Literal, Self, Unpack from codegen.sdk.core.assignment import Assignment from codegen.sdk.core.autocommit import reader, writer @@ -21,6 +21,8 @@ if TYPE_CHECKING: from tree_sitter import Node as TSNode + from codegen.sdk.codebase.flagging.code_flag import CodeFlag + from codegen.sdk.codebase.flagging.enums import FlagKwargs from codegen.sdk.core.detached_symbols.parameter import Parameter from codegen.sdk.core.file import SourceFile from codegen.sdk.core.import_resolution import Import @@ -487,3 +489,26 @@ def convert_to_react_interface(self) -> str | None: if imp.module.source.strip("'").strip('"') in ("react", "prop-types"): imp.remove_if_unused() return interface_name + generic_name + + @writer + def flag(self, **kwargs: Unpack[FlagKwargs]) -> CodeFlag[Self]: + """Flags a TypeScript symbol by adding a flag comment and returning a CodeFlag. + + This implementation first creates the CodeFlag through the standard flagging system, + then adds a TypeScript-specific comment to visually mark the flagged code. + + Args: + **kwargs: Flag keyword arguments including optional 'message' + + Returns: + CodeFlag[Self]: The code flag object for tracking purposes + """ + # First create the standard CodeFlag through the base implementation + code_flag = super().flag(**kwargs) + + # Add a TypeScript comment to visually mark the flag + message = kwargs.get("message", "") + if message: + self.set_inline_comment(f"🚩 {message}") + + return code_flag diff --git a/tests/unit/codegen/sdk/python/function/test_function_flag.py b/tests/unit/codegen/sdk/python/function/test_function_flag.py new file mode 100644 index 000000000..ccc340803 --- /dev/null +++ b/tests/unit/codegen/sdk/python/function/test_function_flag.py @@ -0,0 +1,22 @@ +from codegen.sdk.codebase.factory.get_session import get_codebase_session +from codegen.sdk.enums import ProgrammingLanguage + + +def test_function_flag_with_message(tmpdir): + # language=python + content = """ +def foo(): + pass +""" + with get_codebase_session(tmpdir=tmpdir, files={"test.py": content}, programming_language=ProgrammingLanguage.PYTHON) as codebase: + file = codebase.get_file("test.py") + foo = file.get_function("foo") + + foo.flag(message="This is a test") + codebase.commit() + + expected = """ +def foo(): + pass # 🚩 This is a test +""" + assert file.content == expected diff --git a/tests/unit/codegen/sdk/typescript/function/test_function_flag.py b/tests/unit/codegen/sdk/typescript/function/test_function_flag.py new file mode 100644 index 000000000..1bacd008f --- /dev/null +++ b/tests/unit/codegen/sdk/typescript/function/test_function_flag.py @@ -0,0 +1,24 @@ +from codegen.sdk.codebase.factory.get_session import get_codebase_session +from codegen.sdk.enums import ProgrammingLanguage + + +def test_function_flag_with_message(tmpdir): + # language=typescript + content = """ +function foo() { + return; +} +""" + with get_codebase_session(tmpdir=tmpdir, files={"test.ts": content}, programming_language=ProgrammingLanguage.TYPESCRIPT) as codebase: + file = codebase.get_file("test.ts") + foo = file.get_function("foo") + + foo.flag(message="This is a test") + codebase.commit() + + expected = """ +function foo() { + return; +} // 🚩 This is a test +""" + assert file.content == expected