optimize django debug codemod

pixee · Nov 9, 2023 · 06b0e32 · 06b0e32
1 parent e16e466
commit 06b0e32
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/src/core_codemods/django_debug_flag_on.py b/src/core_codemods/django_debug_flag_on.py
@@ -33,11 +33,11 @@ def rule(cls):
                - settings.py
         """
 
-    def leave_Assign(self, original_node, updated_node):
-        # checks if this module is a settings.py file from django's default directory structure
-        if is_django_settings_file(self.file_context.file_path):
-            return super().leave_Assign(original_node, updated_node)
-        return updated_node
+    def visit_Module(self, _: cst.Module) -> bool:
+        """
+        Only visit module with this codemod if it's a settings.py file.
+        """
+        return is_django_settings_file(self.file_context.file_path)
 
     def on_result_found(self, _, updated_node):
         return updated_node.with_changes(value=cst.Name("False"))