Skip to content

Commit 73fb210

Browse files
kopekCcodegen-team
andauthored
Implements Flag for PYSymbol and TSSymbol (#284)
# Motivation <!-- Why is this change necessary? --> # Content <!-- Please include a summary of the change --> # Testing <!-- How was the change tested? --> # Please check the following before marking your PR as ready for review - [ x] I have added tests for my changes - [ x] I have updated the documentation or added new documentation as needed --------- Co-authored-by: kopekC <[email protected]> Co-authored-by: codegen-team <[email protected]>
1 parent af3033c commit 73fb210

File tree

4 files changed

+98
-2
lines changed

4 files changed

+98
-2
lines changed

src/codegen/sdk/python/symbol.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from __future__ import annotations
22

3-
from typing import TYPE_CHECKING
3+
from typing import TYPE_CHECKING, Self, Unpack
44

55
from codegen.sdk.core.autocommit import reader, writer
66
from codegen.sdk.core.symbol import Symbol
@@ -13,6 +13,8 @@
1313
from tree_sitter import Node as TSNode
1414

1515
from codegen.sdk.codebase.codebase_graph import CodebaseGraph
16+
from codegen.sdk.codebase.flagging.code_flag import CodeFlag
17+
from codegen.sdk.codebase.flagging.enums import FlagKwargs
1618
from codegen.sdk.core.interfaces.has_block import HasBlock
1719
from codegen.sdk.core.node_id_factory import NodeId
1820
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
197199
if auto_format:
198200
comment = " " + PyComment.generate_comment(comment, PyCommentType.SINGLE_LINE)
199201
self.insert_after(comment, fix_indentation=False, newline=False)
202+
203+
@writer
204+
def flag(self, **kwargs: Unpack[FlagKwargs]) -> CodeFlag[Self]:
205+
"""Flags a Python symbol by adding a flag comment and returning a CodeFlag.
206+
207+
This implementation first creates the CodeFlag through the standard flagging system,
208+
then adds a Python-specific comment to visually mark the flagged code.
209+
210+
Args:
211+
**kwargs: Flag keyword arguments including optional 'message'
212+
213+
Returns:
214+
CodeFlag[Self]: The code flag object for tracking purposes
215+
"""
216+
# First create the standard CodeFlag through the base implementation
217+
code_flag = super().flag(**kwargs)
218+
219+
# Add a Python comment to visually mark the flag
220+
message = kwargs.get("message", "")
221+
if message:
222+
self.set_inline_comment(f"🚩 {message}")
223+
224+
return code_flag

src/codegen/sdk/typescript/symbol.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from __future__ import annotations
22

3-
from typing import TYPE_CHECKING, Literal
3+
from typing import TYPE_CHECKING, Literal, Self, Unpack
44

55
from codegen.sdk.core.assignment import Assignment
66
from codegen.sdk.core.autocommit import reader, writer
@@ -21,6 +21,8 @@
2121
if TYPE_CHECKING:
2222
from tree_sitter import Node as TSNode
2323

24+
from codegen.sdk.codebase.flagging.code_flag import CodeFlag
25+
from codegen.sdk.codebase.flagging.enums import FlagKwargs
2426
from codegen.sdk.core.detached_symbols.parameter import Parameter
2527
from codegen.sdk.core.file import SourceFile
2628
from codegen.sdk.core.import_resolution import Import
@@ -487,3 +489,26 @@ def convert_to_react_interface(self) -> str | None:
487489
if imp.module.source.strip("'").strip('"') in ("react", "prop-types"):
488490
imp.remove_if_unused()
489491
return interface_name + generic_name
492+
493+
@writer
494+
def flag(self, **kwargs: Unpack[FlagKwargs]) -> CodeFlag[Self]:
495+
"""Flags a TypeScript symbol by adding a flag comment and returning a CodeFlag.
496+
497+
This implementation first creates the CodeFlag through the standard flagging system,
498+
then adds a TypeScript-specific comment to visually mark the flagged code.
499+
500+
Args:
501+
**kwargs: Flag keyword arguments including optional 'message'
502+
503+
Returns:
504+
CodeFlag[Self]: The code flag object for tracking purposes
505+
"""
506+
# First create the standard CodeFlag through the base implementation
507+
code_flag = super().flag(**kwargs)
508+
509+
# Add a TypeScript comment to visually mark the flag
510+
message = kwargs.get("message", "")
511+
if message:
512+
self.set_inline_comment(f"🚩 {message}")
513+
514+
return code_flag
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
from codegen.sdk.codebase.factory.get_session import get_codebase_session
2+
from codegen.sdk.enums import ProgrammingLanguage
3+
4+
5+
def test_function_flag_with_message(tmpdir):
6+
# language=python
7+
content = """
8+
def foo():
9+
pass
10+
"""
11+
with get_codebase_session(tmpdir=tmpdir, files={"test.py": content}, programming_language=ProgrammingLanguage.PYTHON) as codebase:
12+
file = codebase.get_file("test.py")
13+
foo = file.get_function("foo")
14+
15+
foo.flag(message="This is a test")
16+
codebase.commit()
17+
18+
expected = """
19+
def foo():
20+
pass # 🚩 This is a test
21+
"""
22+
assert file.content == expected
Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
from codegen.sdk.codebase.factory.get_session import get_codebase_session
2+
from codegen.sdk.enums import ProgrammingLanguage
3+
4+
5+
def test_function_flag_with_message(tmpdir):
6+
# language=typescript
7+
content = """
8+
function foo() {
9+
return;
10+
}
11+
"""
12+
with get_codebase_session(tmpdir=tmpdir, files={"test.ts": content}, programming_language=ProgrammingLanguage.TYPESCRIPT) as codebase:
13+
file = codebase.get_file("test.ts")
14+
foo = file.get_function("foo")
15+
16+
foo.flag(message="This is a test")
17+
codebase.commit()
18+
19+
expected = """
20+
function foo() {
21+
return;
22+
} // 🚩 This is a test
23+
"""
24+
assert file.content == expected

0 commit comments

Comments
 (0)