diff --git a/src/codegen/sdk/python/assignment.py b/src/codegen/sdk/python/assignment.py index a85f57a82..2614b6d43 100644 --- a/src/codegen/sdk/python/assignment.py +++ b/src/codegen/sdk/python/assignment.py @@ -12,6 +12,7 @@ from codegen.sdk.python.symbol import PySymbol from codegen.sdk.python.symbol_groups.comment_group import PyCommentGroup from codegen.shared.decorators.docs import noapidoc, py_apidoc +from codegen.shared.logging.get_logger import get_logger if TYPE_CHECKING: from tree_sitter import Node as TSNode @@ -20,6 +21,8 @@ from codegen.sdk.core.node_id_factory import NodeId from codegen.sdk.python.statements.assignment_statement import PyAssignmentStatement +logger = get_logger(__name__) + @py_apidoc class PyAssignment(Assignment["PyAssignmentStatement"], PySymbol): @@ -152,6 +155,9 @@ def remove(self, delete_formatting: bool = True, priority: int = 0, dedupe: bool else: self.parent._values_scheduled_for_removal = [] else: + if name.source == "_": + logger.warning("Attempting to remove '_' in unpacking, command will be ignored. If you wish to remove the statement, remove the other remaining variable(s)!") + return transaction_count = self._active_transactions_on_assignment_names(TransactionPriority.Edit) throwaway = [asgnmt.name == "_" for asgnmt in self.parent.assignments].count(True) # Only edit if we didn't already omit all the other assignments, otherwise just remove the whole thing diff --git a/tests/unit/codegen/sdk/python/expressions/test_unpacking.py b/tests/unit/codegen/sdk/python/expressions/test_unpacking.py index cdf853e37..274329bea 100644 --- a/tests/unit/codegen/sdk/python/expressions/test_unpacking.py +++ b/tests/unit/codegen/sdk/python/expressions/test_unpacking.py @@ -1,5 +1,11 @@ +from typing import TYPE_CHECKING +from unittest.mock import patch + from codegen.sdk.codebase.factory.get_session import get_codebase_session +if TYPE_CHECKING: + from codegen.sdk.core.file import SourceFile + def test_remove_unpacking_assignment(tmpdir) -> None: # language=python @@ -155,3 +161,27 @@ def test_remove_unpacking_assignment_num(tmpdir) -> None: assert len(file2.symbols) == 0 assert file2.source == """""" + + +@patch("codegen.sdk.python.assignment.logger") +def test_unpacking_function_with_underscore_removal(mock_logger, tmpdir: str) -> None: + # language=python + content1 = """ + args, _ = parser.parse_known_args() ##args gets deleted + with open(args.template_path) as f: + print('test') + """ + with get_codebase_session( + tmpdir=tmpdir, + files={ + "file1.py": content1, + }, + ) as codebase: + file1: SourceFile = codebase.get_file("file1.py") + + for symbol in codebase.symbols: + if not symbol.usages: + symbol.remove() + codebase.commit() + assert len(file1.symbols) != 0 + assert mock_logger.warning.call_count == 1