|
| 1 | +/* |
| 2 | + * Copyright (c) Facebook, Inc. and its affiliates. |
| 3 | + * |
| 4 | + * This software may be used and distributed according to the terms of the |
| 5 | + * GNU General Public License version 2. |
| 6 | + */ |
| 7 | + |
| 8 | +use cpython::*; |
| 9 | +use cpython_ext::PyNone; |
| 10 | +use cpython_ext::PyPath; |
| 11 | +use cpython_ext::ResultPyErrExt; |
| 12 | +use spawn_ext::CommandExt; |
| 13 | +use std::cell::RefCell; |
| 14 | +use std::fs; |
| 15 | +use std::io; |
| 16 | +use std::process::Child as RustChild; |
| 17 | +use std::process::Command as RustCommand; |
| 18 | +use std::process::ExitStatus as RustExitStatus; |
| 19 | +use std::process::Stdio as RustStdio; |
| 20 | + |
| 21 | +py_class!(class Command |py| { |
| 22 | + data inner: RefCell<RustCommand>; |
| 23 | + |
| 24 | + /// Constructs a new Command for launching the program at path program, with |
| 25 | + /// the following default configuration: |
| 26 | + /// - No arguments to the program |
| 27 | + /// - Inherit the current process's environment |
| 28 | + /// - Inherit the current process's working directory |
| 29 | + /// - Inherit stdin/stdout/stderr for spawn or status, but create pipes for output |
| 30 | + @staticmethod |
| 31 | + def new(program: String) -> PyResult<Self> { |
| 32 | + let command = RustCommand::new(program); |
| 33 | + Self::create_instance(py, RefCell::new(command)) |
| 34 | + } |
| 35 | + |
| 36 | + /// Adds an argument to pass to the program. |
| 37 | + def arg(&self, arg: &str) -> PyResult<Self> { |
| 38 | + self.mutate_then_clone(py, |c| c.arg(arg)) |
| 39 | + } |
| 40 | + |
| 41 | + /// Adds multiple arguments to pass to the program. |
| 42 | + def args(&self, args: Vec<String>) -> PyResult<Self> { |
| 43 | + self.mutate_then_clone(py, |c| c.args(args)) |
| 44 | + } |
| 45 | + |
| 46 | + /// Inserts or updates an environment variable mapping. |
| 47 | + def env(&self, key: &str, val: &str) -> PyResult<Self> { |
| 48 | + self.mutate_then_clone(py, |c| c.env(key, val)) |
| 49 | + } |
| 50 | + |
| 51 | + /// Adds or updates multiple environment variable mappings. |
| 52 | + def envs(&self, items: Vec<(String, String)>) -> PyResult<Self> { |
| 53 | + self.mutate_then_clone(py, |c| c.envs(items)) |
| 54 | + } |
| 55 | + |
| 56 | + /// Clears the entire environment map for the child process. |
| 57 | + def envclear(&self) -> PyResult<Self> { |
| 58 | + self.mutate_then_clone(py, |c| c.env_clear()) |
| 59 | + } |
| 60 | + |
| 61 | + /// Sets the working directory for the child process. |
| 62 | + def currentdir(&self, dir: &PyPath) -> PyResult<Self> { |
| 63 | + self.mutate_then_clone(py, |c| c.current_dir(dir)) |
| 64 | + } |
| 65 | + |
| 66 | + /// Configuration for the child process's standard input (stdin) handle. |
| 67 | + def stdin(&self, cfg: Stdio) -> PyResult<Self> { |
| 68 | + let f = cfg.to_rust(py).map_pyerr(py)?; |
| 69 | + self.mutate_then_clone(py, |c| c.stdin(f)) |
| 70 | + } |
| 71 | + |
| 72 | + /// Configuration for the child process's standard output (stdout) handle. |
| 73 | + def stdout(&self, cfg: Stdio) -> PyResult<Self> { |
| 74 | + let f = cfg.to_rust(py).map_pyerr(py)?; |
| 75 | + self.mutate_then_clone(py, |c| c.stdout(f)) |
| 76 | + } |
| 77 | + |
| 78 | + /// Configuration for the child process's standard error (stderr) handle. |
| 79 | + def stderr(&self, cfg: Stdio) -> PyResult<Self> { |
| 80 | + let f = cfg.to_rust(py).map_pyerr(py)?; |
| 81 | + self.mutate_then_clone(py, |c| c.stderr(f)) |
| 82 | + } |
| 83 | + |
| 84 | + /// Attempt to avoid inheriting file handles. |
| 85 | + /// Call this before setting up redirections. |
| 86 | + def avoidinherithandles(&self) -> PyResult<Self> { |
| 87 | + self.mutate_then_clone(py, |c| c.avoid_inherit_handles()) |
| 88 | + } |
| 89 | + |
| 90 | + /// Use new session or process group. |
| 91 | + /// Call this after avoidinherithandles. |
| 92 | + def newsession(&self) -> PyResult<Self> { |
| 93 | + self.mutate_then_clone(py, |c| c.new_session()) |
| 94 | + } |
| 95 | + |
| 96 | + /// Executes the command as a child process, returning a handle to it. |
| 97 | + def spawn(&self) -> PyResult<Child> { |
| 98 | + // This is safer than `os.fork()` in Python because Python cannot |
| 99 | + // interrupt between `fork()` and `exec()` due to Rust holding GIL. |
| 100 | + let mut inner = self.inner(py).borrow_mut(); |
| 101 | + let child = inner.spawn().map_pyerr(py)?; |
| 102 | + Child::from_rust(py, child) |
| 103 | + } |
| 104 | + |
| 105 | + /// Spawn the process then forget about it. |
| 106 | + /// File handles are not inherited. stdio will be redirected to /dev/null. |
| 107 | + def spawndetached(&self) -> PyResult<Child> { |
| 108 | + // This is safer than `os.fork()` in Python because Python cannot |
| 109 | + // interrupt between `fork()` and `exec()` due to Rust holding GIL. |
| 110 | + let mut inner = self.inner(py).borrow_mut(); |
| 111 | + let child = inner.spawn_detached().map_pyerr(py)?; |
| 112 | + Child::from_rust(py, child) |
| 113 | + } |
| 114 | + |
| 115 | +}); |
| 116 | + |
| 117 | +impl Command { |
| 118 | + /// Make changes to `inner`, then clone self. |
| 119 | + fn mutate_then_clone( |
| 120 | + &self, |
| 121 | + py: Python, |
| 122 | + func: impl FnOnce(&mut RustCommand) -> &mut RustCommand, |
| 123 | + ) -> PyResult<Self> { |
| 124 | + let mut inner = self.inner(py).borrow_mut(); |
| 125 | + func(&mut inner); |
| 126 | + Ok(self.clone_ref(py)) |
| 127 | + } |
| 128 | +} |
| 129 | + |
| 130 | +py_class!(class Stdio |py| { |
| 131 | + data inner: Box<dyn Fn() -> io::Result<RustStdio> + Send + 'static> ; |
| 132 | + |
| 133 | + /// A new pipe should be arranged to connect the parent and child processes. |
| 134 | + @staticmethod |
| 135 | + def piped() -> PyResult<Self> { |
| 136 | + Self::create_instance(py, Box::new(|| Ok(RustStdio::piped()))) |
| 137 | + } |
| 138 | + |
| 139 | + /// The child inherits from the corresponding parent descriptor. |
| 140 | + @staticmethod |
| 141 | + def inherit() -> PyResult<Self> { |
| 142 | + Self::create_instance(py, Box::new(|| Ok(RustStdio::inherit()))) |
| 143 | + } |
| 144 | + |
| 145 | + /// This stream will be ignored. This is the equivalent of attaching the |
| 146 | + /// stream to /dev/null. |
| 147 | + @staticmethod |
| 148 | + def null() -> PyResult<Self> { |
| 149 | + Self::create_instance(py, Box::new(|| Ok(RustStdio::null()))) |
| 150 | + } |
| 151 | + |
| 152 | + /// Open a file as `Stdio`. |
| 153 | + @staticmethod |
| 154 | + def open(path: &PyPath, read: bool = false, write: bool = false, create: bool = false, append: bool = false) -> PyResult<Self> { |
| 155 | + let path = path.to_path_buf(); |
| 156 | + Self::create_instance(py, Box::new(move || { |
| 157 | + let mut opts = fs::OpenOptions::new(); |
| 158 | + let file = opts.write(write).read(read).create(create).append(append).open(&path)?; |
| 159 | + Ok(file.into()) |
| 160 | + })) |
| 161 | + } |
| 162 | +}); |
| 163 | + |
| 164 | +impl Stdio { |
| 165 | + fn to_rust(&self, py: Python) -> io::Result<RustStdio> { |
| 166 | + self.inner(py)() |
| 167 | + } |
| 168 | +} |
| 169 | + |
| 170 | +py_class!(class Child |py| { |
| 171 | + data inner: RefCell<RustChild>; |
| 172 | + |
| 173 | + /// Forces the child process to exit. If the child has already exited, an |
| 174 | + /// InvalidInput error is returned. |
| 175 | + def kill(&self) -> PyResult<PyNone> { |
| 176 | + let mut inner = self.inner(py).borrow_mut(); |
| 177 | + inner.kill().map_pyerr(py)?; |
| 178 | + Ok(PyNone) |
| 179 | + } |
| 180 | + |
| 181 | + /// Returns the OS-assigned process identifier associated with this child. |
| 182 | + def id(&self) -> PyResult<u32> { |
| 183 | + let inner = self.inner(py).borrow(); |
| 184 | + Ok(inner.id()) |
| 185 | + } |
| 186 | + |
| 187 | + /// Waits for the child to exit completely, returning the status that it |
| 188 | + /// exited with. This function will continue to have the same return value |
| 189 | + /// after it has been called at least once. |
| 190 | + def wait(&self) -> PyResult<ExitStatus> { |
| 191 | + let mut inner = self.inner(py).borrow_mut(); |
| 192 | + let status = inner.wait().map_pyerr(py)?; |
| 193 | + ExitStatus::from_rust(py, status) |
| 194 | + } |
| 195 | + |
| 196 | + /// Attempts to collect the exit status of the child if it has already exited. |
| 197 | + def try_wait(&self) -> PyResult<Option<ExitStatus>> { |
| 198 | + let mut inner = self.inner(py).borrow_mut(); |
| 199 | + match inner.try_wait().map_pyerr(py)? { |
| 200 | + Some(s) => Ok(Some(ExitStatus::from_rust(py, s)?)), |
| 201 | + None => Ok(None) |
| 202 | + } |
| 203 | + } |
| 204 | +}); |
| 205 | + |
| 206 | +impl Child { |
| 207 | + fn from_rust(py: Python, child: RustChild) -> PyResult<Self> { |
| 208 | + Self::create_instance(py, RefCell::new(child)) |
| 209 | + } |
| 210 | +} |
| 211 | + |
| 212 | +py_class!(class ExitStatus |py| { |
| 213 | + data inner: RustExitStatus; |
| 214 | + |
| 215 | + /// Was termination successful? Signal termination is not considered a |
| 216 | + /// success, and success is defined as a zero exit status. |
| 217 | + def success(&self) -> PyResult<bool> { |
| 218 | + Ok(self.inner(py).success()) |
| 219 | + } |
| 220 | + |
| 221 | + /// Returns the exit code of the process, if any. |
| 222 | + /// On Unix, this will return None if the process was terminated by a signal. |
| 223 | + def code(&self) -> PyResult<Option<i32>> { |
| 224 | + Ok(self.inner(py).code()) |
| 225 | + } |
| 226 | +}); |
| 227 | + |
| 228 | +impl ExitStatus { |
| 229 | + fn from_rust(py: Python, status: RustExitStatus) -> PyResult<Self> { |
| 230 | + Self::create_instance(py, status) |
| 231 | + } |
| 232 | +} |
| 233 | + |
| 234 | +pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> { |
| 235 | + let name = [package, "process"].join("."); |
| 236 | + let m = PyModule::new(py, &name)?; |
| 237 | + m.add_class::<Child>(py)?; |
| 238 | + m.add_class::<Command>(py)?; |
| 239 | + m.add_class::<Stdio>(py)?; |
| 240 | + Ok(m) |
| 241 | +} |
0 commit comments