Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/build/lib/airflow/utils/python_virtualenv.py: 30%

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

19from __future__ import annotations 

20 

21import os 

22import sys 

23import warnings 

24 

25import jinja2 

26 

27from airflow.utils.decorators import remove_task_decorator as _remove_task_decorator 

28from airflow.utils.process_utils import execute_in_subprocess 

29 

30 

31def _generate_virtualenv_cmd(tmp_dir: str, python_bin: str, system_site_packages: bool) -> list[str]: 

32 cmd = [sys.executable, "-m", "virtualenv", tmp_dir] 

33 if system_site_packages: 

34 cmd.append("--system-site-packages") 

35 if python_bin is not None: 

36 cmd.append(f"--python={python_bin}") 

37 return cmd 

38 

39 

40def _generate_pip_install_cmd_from_file( 

41 tmp_dir: str, requirements_file_path: str, pip_install_options: list[str] 

42) -> list[str]: 

43 cmd = [f"{tmp_dir}/bin/pip", "install"] + pip_install_options + ["-r"] 

44 return cmd + [requirements_file_path] 

45 

46 

47def _generate_pip_install_cmd_from_list( 

48 tmp_dir: str, requirements: list[str], pip_install_options: list[str] 

49) -> list[str]: 

50 cmd = [f"{tmp_dir}/bin/pip", "install"] + pip_install_options 

51 return cmd + requirements 

52 

53 

54def remove_task_decorator(python_source: str, task_decorator_name: str) -> str: 

55 warnings.warn( 

56 "Import remove_task_decorator from airflow.utils.decorators instead", 

57 DeprecationWarning, 

58 stacklevel=2, 

59 ) 

60 return _remove_task_decorator(python_source, task_decorator_name) 

61 

62 

63def prepare_virtualenv( 

64 venv_directory: str, 

65 python_bin: str, 

66 system_site_packages: bool, 

67 requirements: list[str] | None = None, 

68 requirements_file_path: str | None = None, 

69 pip_install_options: list[str] | None = None, 

70) -> str: 

71 """Creates a virtual environment and installs the additional python packages. 

72 

73 :param venv_directory: The path for directory where the environment will be created. 

74 :param python_bin: Path to the Python executable. 

75 :param system_site_packages: Whether to include system_site_packages in your virtualenv. 

76 See virtualenv documentation for more information. 

77 :param requirements: List of additional python packages. 

78 :param requirements_file_path: Path to the ``requirements.txt`` file. 

79 :return: Path to a binary file with Python in a virtual environment. 

80 """ 

81 if pip_install_options is None: 

82 pip_install_options = [] 

83 

84 virtualenv_cmd = _generate_virtualenv_cmd(venv_directory, python_bin, system_site_packages) 

85 execute_in_subprocess(virtualenv_cmd) 

86 

87 if requirements is not None and requirements_file_path is not None: 

88 raise Exception("Either requirements OR requirements_file_path has to be passed, but not both") 

89 

90 pip_cmd = None 

91 if requirements is not None and len(requirements) != 0: 

92 pip_cmd = _generate_pip_install_cmd_from_list(venv_directory, requirements, pip_install_options) 

93 if requirements_file_path is not None and requirements_file_path: 

94 pip_cmd = _generate_pip_install_cmd_from_file( 

95 venv_directory, requirements_file_path, pip_install_options 

96 ) 

97 

98 if pip_cmd: 

99 execute_in_subprocess(pip_cmd) 

100 

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

102 

103 

104def write_python_script( 

105 jinja_context: dict, 

106 filename: str, 

107 render_template_as_native_obj: bool = False, 

108): 

109 """ 

110 Renders the python script to a file to execute in the virtual environment. 

111 

112 :param jinja_context: The jinja context variables to unpack and replace with its placeholders in the 

113 template file. 

114 :param filename: The name of the file to dump the rendered script to. 

115 :param render_template_as_native_obj: If ``True``, rendered Jinja template would be converted 

116 to a native Python object 

117 """ 

118 template_loader = jinja2.FileSystemLoader(searchpath=os.path.dirname(__file__)) 

119 template_env: jinja2.Environment 

120 if render_template_as_native_obj: 

121 template_env = jinja2.nativetypes.NativeEnvironment( 

122 loader=template_loader, undefined=jinja2.StrictUndefined 

123 ) 

124 else: 

125 template_env = jinja2.Environment(loader=template_loader, undefined=jinja2.StrictUndefined) 

126 template = template_env.get_template("python_virtualenv_script.jinja2") 

127 template.stream(**jinja_context).dump(filename)