3636# --- Helpers ------------------------------------------------------------------
3737
3838
39- def __run_command (command : str , capture : bool = True , sudo_password : str | None = None ) -> CompletedProcess [bytes ]:
39+ def __run_command (
40+ command : str , capture : bool = True , sudo_password : str | None = None , env : dict [str , str ] | None = None
41+ ) -> CompletedProcess [bytes ]:
4042 """Execute a command as a subprocess.
4143 If `sudo_password` is provided and not `None`, the command will be executed with
4244 elevated privileges.
@@ -45,6 +47,7 @@ def __run_command(command: str, capture: bool = True, sudo_password: str | None
4547 :param bool capture: Whether to capture the output of the command.
4648 :param str sudo_password: The password to use when executing the command with
4749 elevated privileges.
50+ :param dict env: The environment variables to use when executing the command.
4851 :return: The result of the command execution.
4952 :rtype: CompletedProcess
5053 """
@@ -58,6 +61,7 @@ def __run_command(command: str, capture: bool = True, sudo_password: str | None
5861 check = True ,
5962 capture_output = capture ,
6063 input = sudo_password .encode () if sudo_password is not None else None ,
64+ env = env ,
6165 )
6266
6367
@@ -80,7 +84,9 @@ def __raise_or_log(exception: CalledProcessError, do_raise: bool) -> None:
8084# --- Public API ---------------------------------------------------------------
8185
8286
83- def execute (command : str , sudo : bool = False , raise_on_error : bool = True ) -> CompletedProcess [bytes ] | None :
87+ def execute (
88+ command : str , sudo : bool = False , raise_on_error : bool = True , env : dict [str , str ] | None = None
89+ ) -> CompletedProcess [bytes ] | None :
8490 """Execute a command in the operating system and wait for it to complete.
8591 Output of the command will be captured and returned after the execution completes.
8692
@@ -97,7 +103,7 @@ def execute(command: str, sudo: bool = False, raise_on_error: bool = True) -> Co
97103 """
98104 try :
99105 logger .debug (f"Running process: { shlex .quote (command )} " )
100- process_result = __run_command (command )
106+ process_result = __run_command (command , env = env )
101107 except CalledProcessError as exception :
102108 # If already running as root, sudo will not work
103109 if not sudo or not os .geteuid ():
@@ -112,7 +118,7 @@ def execute(command: str, sudo: bool = False, raise_on_error: bool = True) -> Co
112118 return None
113119
114120 try :
115- process_result = __run_command (command , sudo_password = sudo_password )
121+ process_result = __run_command (command , sudo_password = sudo_password , env = env )
116122 except CalledProcessError as exception :
117123 sudo_password = None
118124 __raise_or_log (exception , raise_on_error )
@@ -121,15 +127,16 @@ def execute(command: str, sudo: bool = False, raise_on_error: bool = True) -> Co
121127 return process_result
122128
123129
124- def run (command : str ) -> CompletedProcess :
130+ def run (command : str , env : dict [ str , str ] | None = None ) -> CompletedProcess :
125131 """Execute a command in the operating system and wait for it to complete.
126132 Output of the command will not be captured and will be printed to the console
127133 in real-time.
128134
129135 :param str command: The command to execute.
136+ :param dict env: The environment variables to use when executing the command.
130137 """
131138 logger .debug (f"Running process: { shlex .quote (command )} " )
132- return __run_command (command , capture = False )
139+ return __run_command (command , capture = False , env = env )
133140
134141
135142def detached (command : str ) -> Popen [bytes ]:
@@ -141,15 +148,16 @@ def detached(command: str) -> Popen[bytes]:
141148 return Popen (command , shell = True , start_new_session = True , stdout = DEVNULL , stderr = DEVNULL ) # noqa: S602 - intentional use of shell=True
142149
143150
144- def stream (command : str ) -> Generator [str , None , None ]: # noqa: PLR0912
151+ def stream (command : str , env : dict [ str , str ] | None = None ) -> Generator [str , None , None ]: # noqa: PLR0912
145152 """Execute a command in the operating system and stream its output line by line.
146153 :param str command: The command to execute.
154+ :param dict env: The environment variables to use when executing the command.
147155 """
148156 logger .debug (f"Streaming process: { shlex .quote (command )} " )
149157
150158 if not sys .stdin .isatty ():
151159 logger .warning ("STDIN is not a TTY, running command in non-interactive mode" )
152- exec_process = execute (command )
160+ exec_process = execute (command , env = env )
153161
154162 if not exec_process :
155163 yield ""
@@ -164,13 +172,14 @@ def stream(command: str) -> Generator[str, None, None]: # noqa: PLR0912
164172 master , slave = pty .openpty ()
165173
166174 try :
167- process = Popen ( # noqa: S603
168- shlex . split ( command ) ,
175+ process = Popen ( # noqa: S602
176+ command ,
169177 stdout = slave ,
170178 stderr = slave ,
171179 stdin = slave ,
172180 start_new_session = True ,
173- universal_newlines = True ,
181+ shell = True ,
182+ env = env ,
174183 )
175184
176185 received_buffer : bytes = b""
0 commit comments