Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1#
2# Licensed to the Apache Software Foundation (ASF) under one
3# or more contributor license agreements. See the NOTICE file
4# distributed with this work for additional information
5# regarding copyright ownership. The ASF licenses this file
6# to you under the Apache License, Version 2.0 (the
7# "License"); you may not use this file except in compliance
8# with the License. You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing,
13# software distributed under the License is distributed on an
14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15# KIND, either express or implied. See the License for the
16# specific language governing permissions and limitations
17# under the License.
18"""Utilities for creating a virtual environment."""
20from __future__ import annotations
22import os
23import sys
24import warnings
25from pathlib import Path
27import jinja2
28from jinja2 import select_autoescape
30from airflow.utils.decorators import remove_task_decorator as _remove_task_decorator
31from airflow.utils.process_utils import execute_in_subprocess
34def _generate_virtualenv_cmd(tmp_dir: str, python_bin: str, system_site_packages: bool) -> list[str]:
35 cmd = [sys.executable, "-m", "virtualenv", tmp_dir]
36 if system_site_packages:
37 cmd.append("--system-site-packages")
38 if python_bin is not None:
39 cmd.append(f"--python={python_bin}")
40 return cmd
43def _generate_pip_install_cmd_from_file(
44 tmp_dir: str, requirements_file_path: str, pip_install_options: list[str]
45) -> list[str]:
46 return [f"{tmp_dir}/bin/pip", "install", *pip_install_options, "-r", requirements_file_path]
49def _generate_pip_install_cmd_from_list(
50 tmp_dir: str, requirements: list[str], pip_install_options: list[str]
51) -> list[str]:
52 return [f"{tmp_dir}/bin/pip", "install", *pip_install_options, *requirements]
55def _generate_pip_conf(conf_file: Path, index_urls: list[str]) -> None:
56 if index_urls:
57 pip_conf_options = f"index-url = {index_urls[0]}"
58 if len(index_urls) > 1:
59 pip_conf_options += f"\nextra-index-url = {' '.join(x for x in index_urls[1:])}"
60 else:
61 pip_conf_options = "no-index = true"
62 conf_file.write_text(f"[global]\n{pip_conf_options}")
65def remove_task_decorator(python_source: str, task_decorator_name: str) -> str:
66 warnings.warn(
67 "Import remove_task_decorator from airflow.utils.decorators instead",
68 DeprecationWarning,
69 stacklevel=2,
70 )
71 return _remove_task_decorator(python_source, task_decorator_name)
74def prepare_virtualenv(
75 venv_directory: str,
76 python_bin: str,
77 system_site_packages: bool,
78 requirements: list[str] | None = None,
79 requirements_file_path: str | None = None,
80 pip_install_options: list[str] | None = None,
81 index_urls: list[str] | None = None,
82) -> str:
83 """
84 Create a virtual environment and install the additional python packages.
86 :param venv_directory: The path for directory where the environment will be created.
87 :param python_bin: Path to the Python executable.
88 :param system_site_packages: Whether to include system_site_packages in your virtualenv.
89 See virtualenv documentation for more information.
90 :param requirements: List of additional python packages.
91 :param requirements_file_path: Path to the ``requirements.txt`` file.
92 :param pip_install_options: a list of pip install options when installing requirements
93 See 'pip install -h' for available options
94 :param index_urls: an optional list of index urls to load Python packages from.
95 If not provided the system pip conf will be used to source packages from.
96 :return: Path to a binary file with Python in a virtual environment.
97 """
98 if pip_install_options is None:
99 pip_install_options = []
101 if index_urls is not None:
102 _generate_pip_conf(Path(venv_directory) / "pip.conf", index_urls)
104 virtualenv_cmd = _generate_virtualenv_cmd(venv_directory, python_bin, system_site_packages)
105 execute_in_subprocess(virtualenv_cmd)
107 if requirements is not None and requirements_file_path is not None:
108 raise ValueError("Either requirements OR requirements_file_path has to be passed, but not both")
110 pip_cmd = None
111 if requirements is not None and len(requirements) != 0:
112 pip_cmd = _generate_pip_install_cmd_from_list(venv_directory, requirements, pip_install_options)
113 if requirements_file_path is not None and requirements_file_path:
114 pip_cmd = _generate_pip_install_cmd_from_file(
115 venv_directory, requirements_file_path, pip_install_options
116 )
118 if pip_cmd:
119 execute_in_subprocess(pip_cmd)
121 return f"{venv_directory}/bin/python"
124def write_python_script(
125 jinja_context: dict,
126 filename: str,
127 render_template_as_native_obj: bool = False,
128):
129 """
130 Render the python script to a file to execute in the virtual environment.
132 :param jinja_context: The jinja context variables to unpack and replace with its placeholders in the
133 template file.
134 :param filename: The name of the file to dump the rendered script to.
135 :param render_template_as_native_obj: If ``True``, rendered Jinja template would be converted
136 to a native Python object
137 """
138 template_loader = jinja2.FileSystemLoader(searchpath=os.path.dirname(__file__))
139 template_env: jinja2.Environment
140 if render_template_as_native_obj:
141 template_env = jinja2.nativetypes.NativeEnvironment(
142 loader=template_loader, undefined=jinja2.StrictUndefined
143 )
144 else:
145 template_env = jinja2.Environment(
146 loader=template_loader,
147 undefined=jinja2.StrictUndefined,
148 autoescape=select_autoescape(["html", "xml"]),
149 )
150 template = template_env.get_template("python_virtualenv_script.jinja2")
151 template.stream(**jinja_context).dump(filename)