Repository: apache/airflow
Pull request #50446: Preserve all context keys during serialization
Author: login: yannlambret
Labels: name: area:providers
name: provider:standard
Created: 2025-05-10T22:18:15Z | Merged: 2025-05-12T18:00:13Z

When instantiating a PythonVirtualenvOperator, part of the context is propagated based on the requirements constructor parameter. The current logic implementation prevents 'pendulum' or 'apache-airflow' keys to be added to the Airflow context as soon as a version specifier is used.

For instance this requirement list will have the expected behavior:
```python
        op = PythonVirtualenvOperator(
            task_id="task",
            python_callable=func,
            requirements=["pendulum"],
            system_site_packages=False,
        )
```
Whereas this one will cause the `pendulum` objects to be stripped out of the context:
```python
        op = PythonVirtualenvOperator(
            task_id="task",
            python_callable=func,
            requirements=["pendulum<3.0"],
            system_site_packages=False,
        )
```
This PR aims to fix these limitations.


Diff:
diff --git a/providers/standard/src/airflow/providers/standard/operators/python.py b/providers/standard/src/airflow/providers/standard/operators/python.py
index 3d8902dcfc0d1..37aa79c33433e 100644
--- a/providers/standard/src/airflow/providers/standard/operators/python.py
+++ b/providers/standard/src/airflow/providers/standard/operators/python.py
@@ -21,6 +21,7 @@
 import json
 import logging
 import os
+import re
 import shutil
 import subprocess
 import sys
@@ -35,6 +36,9 @@
 from typing import TYPE_CHECKING, Any, Callable, NamedTuple, cast
 
 import lazy_object_proxy
+from packaging.requirements import InvalidRequirement, Requirement
+from packaging.specifiers import InvalidSpecifier
+from packaging.version import InvalidVersion
 
 from airflow.exceptions import (
     AirflowConfigException,
@@ -848,10 +852,38 @@ def execute_callable(self):
 
     def _iter_serializable_context_keys(self):
         yield from self.BASE_SERIALIZABLE_CONTEXT_KEYS
-        if self.system_site_packages or "apache-airflow" in self.requirements:
+
+        found_airflow = found_pendulum = False
+
+        if self.system_site_packages:
+            # If we're using system packages, assume both are present
+            found_airflow = found_pendulum = True
+        else:
+            for raw_str in self.requirements:
+                line = raw_str.strip()
+                # Skip blank lines and full‐line comments
+                if not line or line.startswith("#"):
+                    continue
+
+                # Strip off any inline comment
+                # e.g. turn "foo==1.2.3  # comment" → "foo==1.2.3"
+                req_str = re.sub(r"#.*$", "", line).strip()
+
+                try:
+                    req = Requirement(req_str)
+                except (InvalidRequirement, InvalidSpecifier, InvalidVersion) as e:
+                    raise ValueError(f"Invalid requirement '{raw_str}': {e}") from e
+
+                if req.name == "apache-airflow":
+                    found_airflow = found_pendulum = True
+                    break
+                elif req.name == "pendulum":
+                    found_pendulum = True
+
+        if found_airflow:
             yield from self.AIRFLOW_SERIALIZABLE_CONTEXT_KEYS
             yield from self.PENDULUM_SERIALIZABLE_CONTEXT_KEYS
-        elif "pendulum" in self.requirements:
+        elif found_pendulum:
             yield from self.PENDULUM_SERIALIZABLE_CONTEXT_KEYS
 
 
diff --git a/providers/standard/tests/unit/standard/operators/test_python.py b/providers/standard/tests/unit/standard/operators/test_python.py
index d30938f4f00bf..e8a0091481647 100644
--- a/providers/standard/tests/unit/standard/operators/test_python.py
+++ b/providers/standard/tests/unit/standard/operators/test_python.py
@@ -1451,6 +1451,99 @@ def f(
 
         self.run_as_task(f, serializer=serializer, system_site_packages=False, requirements=None)
 
+    @pytest.mark.parametrize(
+        "requirements, system_site, want_airflow, want_pendulum",
+        [
+            # nothing → just base keys
+            ([], False, False, False),
+            # site-packages → base keys + pendulum keys
+            ([], True, True, True),
+            # apache-airflow / no version constraint
+            (["apache-airflow"], False, True, True),
+            # specific version
+            (["apache-airflow==2.10.2"], False, True, True),
+            # minimum version
+            (["apache-airflow>=2.10"], False, True, True),
+            # pendulum / no version constraint
+            (["pendulum"], False, False, True),
+            # compatible release
+            (["pendulum~=2.1.0"], False, False, True),
+            # other package
+            (["foo==1.0.0"], False, False, False),
+            # with other package
+            (["apache-airflow", "foo"], False, True, True),
+            # full-line comment only
+            (["# comment"], False, False, False),
+            # inline comment after requirement
+            (["apache-airflow==2.10.2  # comment"], False, True, True),
+            # blank line + requirement
+            (["", "pendulum"], False, False, True),
+            # indented comment + requirement
+            (["  # comment", "pendulum~=2.1.0"], False, False, True),
+        ],
+    )
+    def test_iter_serializable_context_keys(self, requirements, system_site, want_airflow, want_pendulum):
+        def func():
+            return "test_return_value"
+
+        op = PythonVirtualenvOperator(
+            task_id="task",
+            python_callable=func,
+            requirements=requirements,
+            system_site_packages=system_site,
+        )
+        keys = set(op._iter_serializable_context_keys())
+
+        base_keys = set(op.BASE_SERIALIZABLE_CONTEXT_KEYS)
+        airflow_keys = set(op.AIRFLOW_SERIALIZABLE_CONTEXT_KEYS)
+        pendulum_keys = set(op.PENDULUM_SERIALIZABLE_CONTEXT_KEYS)
+
+        # BASE keys always present
+        assert base_keys <= keys
+
+        # AIRFLOW keys only when expected
+        if want_airflow:
+            assert airflow_keys <= keys, f"expected AIRFLOW keys for requirements: {requirements}"
+        else:
+            assert not (airflow_keys & keys), f"unexpected AIRFLOW keys for requirements: {requirements}"
+
+        # PENDULUM keys only when expected
+        if want_pendulum:
+            assert pendulum_keys <= keys, f"expected PENDULUM keys for requirements: {requirements}"
+        else:
+            assert not (pendulum_keys & keys), f"unexpected PENDULUM keys for requirements: {requirements}"
+
+    @pytest.mark.parametrize(
+        "invalid_requirement",
+        [
+            # invalid version format
+            "pendulum==3..0",
+            # invalid operator (=< instead of <=)
+            "apache-airflow=<2.0",
+            # same invalid operator on pendulum
+            "pendulum=<3.0",
+            # totally malformed
+            "invalid requirement",
+        ],
+    )
+    def test_iter_serializable_context_keys_invalid_requirement(self, invalid_requirement):
+        def func():
+            return "test_return_value"
+
+        op = PythonVirtualenvOperator(
+            task_id="task",
+            python_callable=func,
+            requirements=[invalid_requirement],
+            system_site_packages=False,
+        )
+
+        with pytest.raises(ValueError) as exc_info:
+            # Consume the generator to trigger parsing
+            list(op._iter_serializable_context_keys())
+
+        msg = str(exc_info.value)
+        assert f"Invalid requirement '{invalid_requirement}'" in msg
+
 
 # when venv tests are run in parallel to other test they create new processes and this might take
 # quite some time in shared docker environment and get some contention even between different containers


Reviews:
user: login: potiuk
body: Nice one - however, maybe instead of regexp you should parse the requirement with packaging requirements and take name of distribution from there: https://packaging.pypa.io/en/stable/requirements.html 

This has additional advantage of actually parsing the requirements and will fail early in case any of the requirements are not following the specification.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-11T08:25:38Z
user: login: potiuk
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-11T22:33:49Z
user: login: potiuk
body: Nice! One, very small nit. Good points with comment removals.
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-11T22:34:27Z
user: login: yannlambret
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-12T06:35:17Z
user: login: shahar1
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-12T17:59:59Z
user: login: amoghrajesh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T05:58:45Z

Review comments:
diff_hunk: @@ -848,10 +852,39 @@ def execute_callable(self):
 
     def _iter_serializable_context_keys(self):
         yield from self.BASE_SERIALIZABLE_CONTEXT_KEYS
-        if self.system_site_packages or "apache-airflow" in self.requirements:
+
+        found_airflow = found_pendulum = False
+
+        if self.system_site_packages:
+            # If we're using system packages, assume both are present
+            found_airflow = found_pendulum = True
+        else:
+            for raw_str in self.requirements:
+                line = raw_str.strip()
+                # Skip blank lines and full‐line comments
+                if not line or line.startswith("#"):
+                    continue
+
+                # Strip off any inline comment
+                # e.g. turn "foo==1.2.3  # comment" → "foo==1.2.3"
+                req_str = re.sub(r"#.*$", "", line).strip()
+
+                try:
+                    req = Requirement(req_str)
+                except (InvalidRequirement, InvalidSpecifier, InvalidVersion) as e:
+                    raise ValueError(f"Invalid requirement '{raw_str}': {e}") from e
+
+                name = req.name.lower()
path: providers/standard/src/airflow/providers/standard/operators/python.py
user: login: potiuk
body: NIT. I think Requirement already normalizes package name - which covers replacement of ".", "_" with "-" and lowercasing the name: 

https://packaging.python.org/en/latest/specifications/name-normalization/
created_at: 2025-05-11T22:33:49Z
author_association: MEMBER
diff_hunk: @@ -848,10 +852,39 @@ def execute_callable(self):
 
     def _iter_serializable_context_keys(self):
         yield from self.BASE_SERIALIZABLE_CONTEXT_KEYS
-        if self.system_site_packages or "apache-airflow" in self.requirements:
+
+        found_airflow = found_pendulum = False
+
+        if self.system_site_packages:
+            # If we're using system packages, assume both are present
+            found_airflow = found_pendulum = True
+        else:
+            for raw_str in self.requirements:
+                line = raw_str.strip()
+                # Skip blank lines and full‐line comments
+                if not line or line.startswith("#"):
+                    continue
+
+                # Strip off any inline comment
+                # e.g. turn "foo==1.2.3  # comment" → "foo==1.2.3"
+                req_str = re.sub(r"#.*$", "", line).strip()
+
+                try:
+                    req = Requirement(req_str)
+                except (InvalidRequirement, InvalidSpecifier, InvalidVersion) as e:
+                    raise ValueError(f"Invalid requirement '{raw_str}': {e}") from e
+
+                name = req.name.lower()
path: providers/standard/src/airflow/providers/standard/operators/python.py
user: login: yannlambret
body: Good point, I'm changing this 👍🏻 
created_at: 2025-05-12T06:35:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1451,6 +1451,99 @@ def f(
 
         self.run_as_task(f, serializer=serializer, system_site_packages=False, requirements=None)
 
+    @pytest.mark.parametrize(
+        "requirements, system_site, want_airflow, want_pendulum",
+        [
+            # nothing → just base keys
+            ([], False, False, False),
+            # site-packages → base keys + pendulum keys
+            ([], True, True, True),
+            # apache-airflow / no version constraint
+            (["apache-airflow"], False, True, True),
+            # specific version
+            (["apache-airflow==2.10.2"], False, True, True),
+            # minimum version
+            (["apache-airflow>=2.10"], False, True, True),
+            # pendulum / no version constraint
+            (["pendulum"], False, False, True),
+            # compatible release
+            (["pendulum~=2.1.0"], False, False, True),
+            # other package
+            (["foo==1.0.0"], False, False, False),
+            # with other package
+            (["apache-airflow", "foo"], False, True, True),
+            # full-line comment only
+            (["# comment"], False, False, False),
+            # inline comment after requirement
+            (["apache-airflow==2.10.2  # comment"], False, True, True),
+            # blank line + requirement
+            (["", "pendulum"], False, False, True),
+            # indented comment + requirement
+            (["  # comment", "pendulum~=2.1.0"], False, False, True),
+        ],
+    )
+    def test_iter_serializable_context_keys(self, requirements, system_site, want_airflow, want_pendulum):
+        def func():
+            return "test_return_value"
+
+        op = PythonVirtualenvOperator(
+            task_id="task",
+            python_callable=func,
+            requirements=requirements,
+            system_site_packages=system_site,
+        )
+        keys = set(op._iter_serializable_context_keys())
+
+        base_keys = set(op.BASE_SERIALIZABLE_CONTEXT_KEYS)
+        airflow_keys = set(op.AIRFLOW_SERIALIZABLE_CONTEXT_KEYS)
+        pendulum_keys = set(op.PENDULUM_SERIALIZABLE_CONTEXT_KEYS)
+
+        # BASE keys always present
+        assert base_keys <= keys
+
+        # AIRFLOW keys only when expected
+        if want_airflow:
+            assert airflow_keys <= keys, f"expected AIRFLOW keys for requirements: {requirements}"
+        else:
+            assert not (airflow_keys & keys), f"unexpected AIRFLOW keys for requirements: {requirements}"
+
+        # PENDULUM keys only when expected
+        if want_pendulum:
+            assert pendulum_keys <= keys, f"expected PENDULUM keys for requirements: {requirements}"
+        else:
+            assert not (pendulum_keys & keys), f"unexpected PENDULUM keys for requirements: {requirements}"
+
+    @pytest.mark.parametrize(
+        "invalid_requirement",
+        [
+            # invalid version format
+            "pendulum==3..0",
+            # invalid operator (=< instead of <=)
+            "apache-airflow=<2.0",
+            # same invalid operator on pendulum
+            "pendulum=<3.0",
+            # totally malformed
+            "invalid requirement",
path: providers/standard/tests/unit/standard/operators/test_python.py
user: login: amoghrajesh
body: We do not account for the case when more than one are passed.
created_at: 2025-05-13T05:58:42Z
author_association: CONTRIBUTOR

PR comments:
user: login: boring-cyborg[bot]
created_at: 2025-05-10T22:18:19Z
body: Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide (https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
Here are some useful points:
- Pay attention to the quality of your code (ruff, mypy and type annotations). Our [pre-commits]( https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst#prerequisites-for-pre-commit-hooks) will help you with that.
- In case of a new feature add useful documentation (in docstrings or in `docs/` directory). Adding a new operator? Check this short [guide](https://github.com/apache/airflow/blob/main/airflow-core/docs/howto/custom-operator.rst) Consider adding an example DAG that shows how users should use it.
- Consider using [Breeze environment](https://github.com/apache/airflow/blob/main/dev/breeze/doc/README.rst) for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
- Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
- Please follow [ASF Code of Conduct](https://www.apache.org/foundation/policies/conduct) for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
- Be sure to read the [Airflow Coding style]( https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#coding-style-and-best-practices).
- Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
Apache Airflow is a community-driven project and together we are making it better 🚀.
In case of doubts contact the developers at:
Mailing List: dev@airflow.apache.org
Slack: https://s.apache.org/airflow-slack

author_association: NONE
user: login: yannlambret
created_at: 2025-05-11T13:14:19Z
body: Hello @potiuk, thank you for the suggestion. I modified the code to use `packaging.requirements`. I did my best to implement a reasonably robust parsing logic, by taking care of a few edge cases (comments, etc.). Let me know what you think. 
author_association: CONTRIBUTOR
user: login: boring-cyborg[bot]
created_at: 2025-05-12T18:00:16Z
body: Awesome work, congrats on your first merged pull request! You are invited to check our [Issue Tracker](https://github.com/apache/airflow/issues) for additional contributions.

author_association: NONE