rename skip_exit_code to skip_on_exit_code and allow providing multiple codes (#30692)

* rename skip_exit_code to skip_on_exit_code and allow providing multiple codes

* replace list type by Container

Author: hussein-awala

airflow/operators/python.py

@@ -26,6 +26,7 @@
 import types
 import warnings
 from abc import ABCMeta, abstractmethod
+from collections.abc import Container
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from textwrap import dedent
@@ -471,7 +472,7 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator):
     :param expect_airflow: expect Airflow to be installed in the target environment. If true, the operator
         will raise warning if Airflow is not installed, and it will attempt to load Airflow
         macros when starting.
-    :param skip_exit_code: If python_callable exits with this exit code, leave the task
+    :param skip_on_exit_code: If python_callable exits with this exit code, leave the task
         in ``skipped`` state (default: None). If set to ``None``, any non-zero
         exit code will be treated as a failure.
     """
@@ -494,7 +495,7 @@ def __init__(
         templates_dict: dict | None = None,
         templates_exts: list[str] | None = None,
         expect_airflow: bool = True,
-        skip_exit_code: int | None = None,
+        skip_on_exit_code: int | Container[int] | None = None,
         **kwargs,
     ):
         if (
@@ -518,7 +519,13 @@ def __init__(
         self.python_version = python_version
         self.system_site_packages = system_site_packages
         self.pip_install_options = pip_install_options
-        self.skip_exit_code = skip_exit_code
+        self.skip_on_exit_code = (
+            skip_on_exit_code
+            if isinstance(skip_on_exit_code, Container)
+            else [skip_on_exit_code]
+            if skip_on_exit_code
+            else []
+        )
         super().__init__(
             python_callable=python_callable,
             use_dill=use_dill,
@@ -557,8 +564,8 @@ def execute_callable(self):
             try:
                 result = self._execute_python_callable_in_subprocess(python_path, tmp_path)
             except subprocess.CalledProcessError as e:
-                if self.skip_exit_code and e.returncode == self.skip_exit_code:
-                    raise AirflowSkipException(f"Process exited with code {self.skip_exit_code}. Skipping.")
+                if e.returncode in self.skip_on_exit_code:
+                    raise AirflowSkipException(f"Process exited with code {e.returncode}. Skipping.")
                 else:
                     raise
             return result

tests/operators/test_python.py

@@ -938,12 +938,15 @@ def test_virtualenv_serializable_context_fields(self, create_task_instance):
         "extra_kwargs, actual_exit_code, expected_state",
         [
             (None, 99, TaskInstanceState.FAILED),
-            ({"skip_exit_code": 100}, 100, TaskInstanceState.SKIPPED),
-            ({"skip_exit_code": 100}, 101, TaskInstanceState.FAILED),
-            ({"skip_exit_code": None}, 0, TaskInstanceState.SUCCESS),
+            ({"skip_on_exit_code": 100}, 100, TaskInstanceState.SKIPPED),
+            ({"skip_on_exit_code": [100]}, 100, TaskInstanceState.SKIPPED),
+            ({"skip_on_exit_code": (100, 101)}, 100, TaskInstanceState.SKIPPED),
+            ({"skip_on_exit_code": 100}, 101, TaskInstanceState.FAILED),
+            ({"skip_on_exit_code": [100, 102]}, 101, TaskInstanceState.FAILED),
+            ({"skip_on_exit_code": None}, 0, TaskInstanceState.SUCCESS),
         ],
     )
-    def test_skip_exit_code(self, extra_kwargs, actual_exit_code, expected_state):
+    def test_on_skip_exit_code(self, extra_kwargs, actual_exit_code, expected_state):
         def f(exit_code):
             if exit_code != 0:
                 raise SystemExit(exit_code)