Handle 'control' mode for treated flag methods

Author: ffmmjj

piranha_python/codemods.py

@@ -35,12 +35,21 @@ def add_args(arg_parser):  # pragma: no cover
             type=str,
             required=False,
         )
+        arg_parser.add_argument(
+            "--mode",
+            dest="mode",
+            metavar="MODE",
+            help="Execution mode - can be 'treated' or 'control'",
+            type=str,
+            required=False,
+        )
 
-    def __init__(self, context, flag_name, flag_resolution_methods, ignored_module_check_fn_path=None):
+    def __init__(self, context, flag_name, flag_resolution_methods, ignored_module_check_fn_path=None, mode="treated"):
         super().__init__(context)
         self.flag_name = flag_name
         self.is_in_feature_flag_block = False
         self.found_return_stmt_in_ff_block = False
+        self.running_in_treated_mode = mode == "treated"
 
         if ignored_module_check_fn_path is None:
             ignored_module_check_fn_path = self.DEFAULT_TEST_MODULE_CHECK_PATH
@@ -132,18 +141,20 @@ def leave_If(self, original_node, updated_node):
         if not self.is_in_feature_flag_block:
             return updated_node
 
-        if matchers.matches(updated_node.test, self.flag_resolution_matcher):
-            if self.is_treatment_method:
+        if self._should_assume_that_flag_is_true():
+            if matchers.matches(updated_node.test, self.flag_resolution_matcher):
                 replaced_node = updated_node.body
-            else:
-                replaced_node = updated_node.orelse.body
-        elif matchers.matches(updated_node.test, _inside_not_matcher(self.flag_resolution_matcher)):
-            if self.is_treatment_method:
+            elif matchers.matches(updated_node.test, _inside_not_matcher(self.flag_resolution_matcher)):
                 replaced_node = updated_node.orelse.body
             else:
-                replaced_node = updated_node.body
+                return updated_node
         else:
-            return updated_node
+            if matchers.matches(updated_node.test, self.flag_resolution_matcher):
+                replaced_node = updated_node.orelse.body
+            elif matchers.matches(updated_node.test, _inside_not_matcher(self.flag_resolution_matcher)):
+                replaced_node = updated_node.body
+            else:
+                return updated_node
 
         return_statements = matchers.findall(replaced_node, matchers.Return())
         self.found_return_stmt_in_ff_block = len(return_statements) > 0
@@ -196,6 +207,11 @@ def _updated_tuple_assignment(self, updated_node):
             ),
         )
 
+    def _should_assume_that_flag_is_true(self):
+        return (self.is_treatment_method and self.running_in_treated_mode) or (
+            not self.is_treatment_method and not self.running_in_treated_mode
+        )
+
 
 def _matches_flag_name(flag_name, n):
     return matchers.matches(n[0].value, matchers.Name(flag_name))

test/test_codemods.py

@@ -251,6 +251,85 @@ def test_keeps_ELSE_block_body_and_remainder_when_IF_block_contains_return_state
             flag_resolution_methods="is_flag_active",
         )
 
+    # Tests covering "control" mode instead of the default "treated" mode
+    def test_keeps_ELSE_block_body_when_mode_is_control(self):
+        self.assertCodemod(
+            _with_correct_indentation(
+                """\
+            if is_flag_active(%s):
+                print('Flag is active')
+            else:
+                print('This is not related to the feature flag value at all')
+
+            print('Completely unrelated statement')
+            """
+                % FEATURE_FLAG_NAME
+            ),
+            _with_correct_indentation(
+                """\
+            print('This is not related to the feature flag value at all')
+
+            print('Completely unrelated statement')
+            """
+            ),
+            flag_name=FEATURE_FLAG_NAME,
+            flag_resolution_methods="is_flag_active",
+            mode="control",
+        )
+
+    def test_keeps_only_ELSE_block_body_when_it_contains_return_statement_and_mode_is_control(self):
+        self.assertCodemod(
+            _with_correct_indentation(
+                """\
+            if is_flag_active(%s):
+                print('Flag is active')
+            else:
+                print('This is not related to the feature flag value at all')
+                return 0
+
+            print('Completely unrelated statement')
+            """
+                % FEATURE_FLAG_NAME
+            ),
+            _with_correct_indentation(
+                """\
+            print('This is not related to the feature flag value at all')
+            return 0
+            """
+            ),
+            flag_name=FEATURE_FLAG_NAME,
+            flag_resolution_methods="is_flag_active",
+            mode="control",
+        )
+
+    def test_keeps_ELSE_block_body_and_remainder_when_IF_block_contains_return_statement_but_ELSE_doesnt_and_mode_is_control(
+        self
+    ):
+        self.assertCodemod(
+            _with_correct_indentation(
+                """\
+            if is_flag_active(%s):
+                print('Flag is active')
+                return 0
+            else:
+                print('This is not related to the feature flag value at all')
+
+            print('Completely unrelated statement')
+            """
+                % FEATURE_FLAG_NAME
+            ),
+            _with_correct_indentation(
+                """\
+            print('This is not related to the feature flag value at all')
+
+            print('Completely unrelated statement')
+            """
+            ),
+            flag_name=FEATURE_FLAG_NAME,
+            flag_resolution_methods="is_flag_active",
+            mode="control",
+        )
+
     # Tests covering the presence of docstrings in changed code
     def test_keeps_single_lined_docstring_in_module(self):
         self.assertCodemod(