This is an automated email from the ASF dual-hosted git repository. potiuk pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push: new d5a75446a6 Fix import future annotations in venv jinja template (#40208) d5a75446a6 is described below commit d5a75446a62ba8804879177ce394c7456adaa4d3 Author: phi-friday <phi.fri...@gmail.com> AuthorDate: Fri Jun 14 14:04:37 2024 +0900 Fix import future annotations in venv jinja template (#40208) --- airflow/utils/python_virtualenv_script.jinja2 | 1 + tests/decorators/test_python_virtualenv.py | 30 +++++++++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/airflow/utils/python_virtualenv_script.jinja2 b/airflow/utils/python_virtualenv_script.jinja2 index 4199a47130..2ff417985e 100644 --- a/airflow/utils/python_virtualenv_script.jinja2 +++ b/airflow/utils/python_virtualenv_script.jinja2 @@ -16,6 +16,7 @@ specific language governing permissions and limitations under the License. -#} +from __future__ import annotations import {{ pickling_library }} import sys diff --git a/tests/decorators/test_python_virtualenv.py b/tests/decorators/test_python_virtualenv.py index b91bcaae36..0f7ab6918d 100644 --- a/tests/decorators/test_python_virtualenv.py +++ b/tests/decorators/test_python_virtualenv.py @@ -21,12 +21,14 @@ import datetime import sys from importlib.util import find_spec from subprocess import CalledProcessError +from typing import Any import pytest from airflow.decorators import setup, task, teardown from airflow.exceptions import RemovedInAirflow3Warning from airflow.utils import timezone +from airflow.utils.state import TaskInstanceState pytestmark = pytest.mark.db_test @@ -37,6 +39,8 @@ DILL_MARKER = pytest.mark.skipif(not DILL_INSTALLED, reason="`dill` is not insta CLOUDPICKLE_INSTALLED = find_spec("cloudpickle") is not None CLOUDPICKLE_MARKER = pytest.mark.skipif(not CLOUDPICKLE_INSTALLED, reason="`cloudpickle` is not installed") +_Invalid = Any + class TestPythonVirtualenvDecorator: @CLOUDPICKLE_MARKER @@ -350,3 +354,29 @@ class TestPythonVirtualenvDecorator: assert teardown_task.is_teardown assert teardown_task.on_failure_fail_dagrun is on_failure_fail_dagrun ret.operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + + def test_invalid_annotation(self, dag_maker): + import uuid + + unique_id = uuid.uuid4().hex + value = {"unique_id": unique_id} + + # Functions that throw an error + # if `from __future__ import annotations` is missing + @task.virtualenv(multiple_outputs=False, do_xcom_push=True) + def in_venv(value: dict[str, _Invalid]) -> _Invalid: + assert isinstance(value, dict) + return value["unique_id"] + + with dag_maker(): + ret = in_venv(value) + + dr = dag_maker.create_dagrun() + ret.operator.run(start_date=dr.execution_date, end_date=dr.execution_date) + ti = dr.get_task_instances()[0] + + assert ti.state == TaskInstanceState.SUCCESS + + xcom = ti.xcom_pull(task_ids=ti.task_id, key="return_value") + assert isinstance(xcom, str) + assert xcom == unique_id