diff --git a/src/codemodder/codemods/base_codemod.py b/src/codemodder/codemods/base_codemod.py index b83f00c42..6b1b44480 100644 --- a/src/codemodder/codemods/base_codemod.py +++ b/src/codemodder/codemods/base_codemod.py @@ -75,14 +75,17 @@ def node_position(self, node): # See https://github.com/Instagram/LibCST/blob/main/libcst/_metadata_dependent.py#L112 return self.get_metadata(self.METADATA_DEPENDENCIES[0], node) - def add_change(self, node, description): + def add_change(self, node, description: str, start: bool = True): position = self.node_position(node) - self.add_change_from_position(position, description) + self.add_change_from_position(position, description, start) - def add_change_from_position(self, position: CodeRange, description): + def add_change_from_position( + self, position: CodeRange, description: str, start: bool = True + ): + lineno = position.start.line if start else position.end.line self.file_context.codemod_changes.append( Change( - lineNumber=position.start.line, + lineNumber=lineno, description=description, ) ) diff --git a/src/codemodder/codemods/utils.py b/src/codemodder/codemods/utils.py index a4a9c0ab2..c947a999c 100644 --- a/src/codemodder/codemods/utils.py +++ b/src/codemodder/codemods/utils.py @@ -91,3 +91,10 @@ def get_function_name_node(call: cst.Call) -> Optional[cst.Name]: case cst.Attribute(): return call.func.attr return None + + +def is_assigned_to_True(original_node: cst.Assign): + return ( + isinstance(original_node.value, cst.Name) + and original_node.value.value == "True" + ) diff --git a/src/core_codemods/django_session_cookie_secure_off.py b/src/core_codemods/django_session_cookie_secure_off.py index f41de07b6..bf1bf4b72 100644 --- a/src/core_codemods/django_session_cookie_secure_off.py +++ b/src/core_codemods/django_session_cookie_secure_off.py @@ -1,12 +1,7 @@ -from typing import List import libcst as cst -from libcst.codemod import Codemod, CodemodContext -from codemodder.change import Change -from codemodder.codemods.base_visitor import BaseTransformer from codemodder.codemods.api import SemgrepCodemod from codemodder.codemods.base_codemod import ReviewGuidance -from codemodder.codemods.utils import is_django_settings_file -from codemodder.file_context import FileContext +from codemodder.codemods.utils import is_django_settings_file, is_assigned_to_True class DjangoSessionCookieSecureOff(SemgrepCodemod): @@ -65,12 +60,7 @@ def leave_Module( # something else and we changed it in `leave_Assign`. return updated_node - # line_number is the end of the module where we will insert the new flag. - pos_to_match = self.node_position(original_node) - line_number = pos_to_match.end.line - self.file_context.codemod_changes.append( - Change(line_number, self.CHANGE_DESCRIPTION) - ) + self.add_change(original_node, self.CHANGE_DESCRIPTION, start=False) final_line = cst.parse_statement("SESSION_COOKIE_SECURE = True") new_body = updated_node.body + (final_line,) return updated_node.with_changes(body=new_body) @@ -103,10 +93,3 @@ def is_session_cookie_secure(original_node: cst.Assign): return ( isinstance(target_var, cst.Name) and target_var.value == "SESSION_COOKIE_SECURE" ) - - -def is_assigned_to_True(original_node: cst.Assign): - return ( - isinstance(original_node.value, cst.Name) - and original_node.value.value == "True" - )