Implements Flag for PYSymbol and TSSymbol

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

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

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

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