Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/airflow/utils/python_virtualenv.py: 31%

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

55 statements  

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.""" 

19 

20from __future__ import annotations 

21 

22import os 

23import sys 

24import warnings 

25from pathlib import Path 

26 

27import jinja2 

28from jinja2 import select_autoescape 

29 

30from airflow.utils.decorators import remove_task_decorator as _remove_task_decorator 

31from airflow.utils.process_utils import execute_in_subprocess 

32 

33 

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 

41 

42 

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] 

47 

48 

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] 

53 

54 

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}") 

63 

64 

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) 

72 

73 

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. 

85 

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 = [] 

100 

101 if index_urls is not None: 

102 _generate_pip_conf(Path(venv_directory) / "pip.conf", index_urls) 

103 

104 virtualenv_cmd = _generate_virtualenv_cmd(venv_directory, python_bin, system_site_packages) 

105 execute_in_subprocess(virtualenv_cmd) 

106 

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") 

109 

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 ) 

117 

118 if pip_cmd: 

119 execute_in_subprocess(pip_cmd) 

120 

121 return f"{venv_directory}/bin/python" 

122 

123 

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. 

131 

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)