Skip to content

Commit b61ba9d

Browse files
tomcodgentkfoss
andauthored
[CG-10935] fix: issues with unpacking (#724)
# 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 - [ ] I have added tests for my changes - [ ] I have updated the documentation or added new documentation as needed --------- Co-authored-by: tomcodgen <191515280+tomcodgen@users.noreply.github.com>
1 parent cf45221 commit b61ba9d

File tree

2 files changed

+36
-0
lines changed

2 files changed

+36
-0
lines changed

src/codegen/sdk/python/assignment.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from codegen.sdk.python.symbol import PySymbol
1313
from codegen.sdk.python.symbol_groups.comment_group import PyCommentGroup
1414
from codegen.shared.decorators.docs import noapidoc, py_apidoc
15+
from codegen.shared.logging.get_logger import get_logger
1516

1617
if TYPE_CHECKING:
1718
from tree_sitter import Node as TSNode
@@ -20,6 +21,8 @@
2021
from codegen.sdk.core.node_id_factory import NodeId
2122
from codegen.sdk.python.statements.assignment_statement import PyAssignmentStatement
2223

24+
logger = get_logger(__name__)
25+
2326

2427
@py_apidoc
2528
class PyAssignment(Assignment["PyAssignmentStatement"], PySymbol):
@@ -152,6 +155,9 @@ def remove(self, delete_formatting: bool = True, priority: int = 0, dedupe: bool
152155
else:
153156
self.parent._values_scheduled_for_removal = []
154157
else:
158+
if name.source == "_":
159+
logger.warning("Attempting to remove '_' in unpacking, command will be ignored. If you wish to remove the statement, remove the other remaining variable(s)!")
160+
return
155161
transaction_count = self._active_transactions_on_assignment_names(TransactionPriority.Edit)
156162
throwaway = [asgnmt.name == "_" for asgnmt in self.parent.assignments].count(True)
157163
# Only edit if we didn't already omit all the other assignments, otherwise just remove the whole thing

tests/unit/codegen/sdk/python/expressions/test_unpacking.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,11 @@
1+
from typing import TYPE_CHECKING
2+
from unittest.mock import patch
3+
14
from codegen.sdk.codebase.factory.get_session import get_codebase_session
25

6+
if TYPE_CHECKING:
7+
from codegen.sdk.core.file import SourceFile
8+
39

410
def test_remove_unpacking_assignment(tmpdir) -> None:
511
# language=python
@@ -155,3 +161,27 @@ def test_remove_unpacking_assignment_num(tmpdir) -> None:
155161

156162
assert len(file2.symbols) == 0
157163
assert file2.source == """"""
164+
165+
166+
@patch("codegen.sdk.python.assignment.logger")
167+
def test_unpacking_function_with_underscore_removal(mock_logger, tmpdir: str) -> None:
168+
# language=python
169+
content1 = """
170+
args, _ = parser.parse_known_args() ##args gets deleted
171+
with open(args.template_path) as f:
172+
print('test')
173+
"""
174+
with get_codebase_session(
175+
tmpdir=tmpdir,
176+
files={
177+
"file1.py": content1,
178+
},
179+
) as codebase:
180+
file1: SourceFile = codebase.get_file("file1.py")
181+
182+
for symbol in codebase.symbols:
183+
if not symbol.usages:
184+
symbol.remove()
185+
codebase.commit()
186+
assert len(file1.symbols) != 0
187+
assert mock_logger.warning.call_count == 1

0 commit comments

Comments
 (0)