3diy
/
moonraker
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

shell_command: allow optional success codes 

Some commands may return codes other than zero on a successful
return.  Allow callers to specify an optional list of return codes that
will return success without raising an exception.

Signed-off-by:  Eric Callahan <arksine.code@gmail.com>
This commit is contained in:
Eric Callahan 2023-02-07 17:56:07 -05:00
parent c69441955a
commit d9579c9374
No known key found for this signature in database
GPG Key ID: 5A1EB336DFB4C71B
1 changed files with 103 additions and 81 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

76
moonraker/components/shell_command.py
View File

@ -28,7 +28,8 @@ if TYPE_CHECKING:
OutputCallback = Optional[Callable[[bytes], None]]
class ShellCommandError(ServerError):
def __init__(self,
def __init__(
self,
message: str,
return_code: Optional[int],
stdout: Optional[bytes] = b"",
@ -41,7 +42,8 @@ class ShellCommandError(ServerError):
self.return_code = return_code
class ShellCommandProtocol(asyncio.subprocess.SubprocessStreamProtocol):
def __init__(self,
def __init__(
self,
limit: int,
loop: asyncio.events.AbstractEventLoop,
program_name: str = "",
@ -58,8 +60,8 @@ class ShellCommandProtocol(asyncio.subprocess.SubprocessStreamProtocol):
self.log_stderr = log_stderr
self.pending_data: List[bytes] = [b"", b""]
def connection_made(self,
transport: asyncio.transports.BaseTransport
def connection_made(
self, transport: asyncio.transports.BaseTransport
) -> None:
self._transport = transport
assert isinstance(transport, asyncio.SubprocessTransport)
@ -74,10 +76,11 @@ class ShellCommandProtocol(asyncio.subprocess.SubprocessStreamProtocol):
stdin_transport = transport.get_pipe_transport(0)
if stdin_transport is not None:
self.stdin = asyncio.streams.StreamWriter(
stdin_transport,
stdin_transport, # type: ignore
protocol=self,
reader=None,
loop=self._loop)
loop=self._loop
)
def pipe_data_received(self, fd: int, data: bytes | str) -> None:
cb = None
@ -103,9 +106,8 @@ class ShellCommandProtocol(asyncio.subprocess.SubprocessStreamProtocol):
continue
cb(line)
def pipe_connection_lost(self,
fd: int,
exc: Exception | None
def pipe_connection_lost(
self, fd: int, exc: Exception | None
) -> None:
cb = None
pending = b""
@ -124,7 +126,8 @@ class ShellCommand:
IDX_SIGINT = 0
IDX_SIGTERM = 1
IDX_SIGKILL = 2
def __init__(self,
def __init__(
self,
factory: ShellCommandFactory,
cmd: str,
std_out_callback: OutputCallback,
@ -178,12 +181,14 @@ class ShellCommand:
self.return_code = self.proc = None
self.cancelled = False
async def run(self,
async def run(
self,
timeout: float = 2.,
verbose: bool = True,
log_complete: bool = True,
sig_idx: int = 1,
proc_input: Optional[str] = None
proc_input: Optional[str] = None,
success_codes: Optional[List[int]] = None
) -> bool:
async with self.run_lock:
self.factory.add_running_command(self)
@ -217,14 +222,18 @@ class ShellCommand:
else:
complete = not self.cancelled
self.factory.remove_running_command(self)
return self._check_proc_success(complete, log_complete)
return self._check_proc_success(
complete, log_complete, success_codes
)
async def run_with_response(self,
async def run_with_response(
self,
timeout: float = 2.,
retries: int = 1,
log_complete: bool = True,
sig_idx: int = 1,
proc_input: Optional[str] = None
proc_input: Optional[str] = None,
success_codes: Optional[List[int]] = None
) -> str:
async with self.run_lock:
self.factory.add_running_command(self)
@ -252,7 +261,9 @@ class ShellCommand:
logging.info(
f"{self.command[0]}: "
f"{stderr.decode(errors='ignore')}")
if self._check_proc_success(complete, log_complete):
if self._check_proc_success(
complete, log_complete, success_codes
):
self.factory.remove_running_command(self)
return stdout.decode(errors='ignore').rstrip("\n")
if stdout:
@ -268,7 +279,8 @@ class ShellCommand:
f"Error running shell command: '{self.name}'",
self.return_code, stdout, stderr)
async def _create_subprocess(self,
async def _create_subprocess(
self,
use_callbacks: bool = False,
has_input: bool = False
) -> bool:
@ -305,13 +317,17 @@ class ShellCommand:
return False
return True
def _check_proc_success(self,
def _check_proc_success(
self,
complete: bool,
log_complete: bool
log_complete: bool,
success_codes: Optional[List[int]] = None
) -> bool:
assert self.proc is not None
if success_codes is None:
success_codes = [0]
self.return_code = self.proc.returncode
success = self.return_code == 0 and complete
success = self.return_code in success_codes and complete
if success:
msg = f"Command ({self.name}) successfully finished"
elif self.cancelled:
@ -339,7 +355,8 @@ class ShellCommandFactory:
except KeyError:
pass
def build_shell_command(self,
def build_shell_command(
self,
cmd: str,
callback: OutputCallback = None,
std_err_callback: OutputCallback = None,
@ -347,10 +364,12 @@ class ShellCommandFactory:
log_stderr: bool = False,
cwd: Optional[str] = None
) -> ShellCommand:
return ShellCommand(self, cmd, callback, std_err_callback, env,
log_stderr, cwd)
return ShellCommand(
self, cmd, callback, std_err_callback, env, log_stderr, cwd
)
def exec_cmd(self,
def exec_cmd(
self,
cmd: str,
timeout: float = 2.,
retries: int = 1,
@ -359,12 +378,15 @@ class ShellCommandFactory:
log_complete: bool = True,
log_stderr: bool = False,
env: Optional[Dict[str, str]] = None,
cwd: Optional[str] = None
cwd: Optional[str] = None,
success_codes: Optional[List[int]] = None
) -> Awaitable:
scmd = ShellCommand(self, cmd, None, None, env,
log_stderr, cwd)
coro = scmd.run_with_response(timeout, retries, log_complete,
sig_idx, proc_input)
coro = scmd.run_with_response(
timeout, retries, log_complete, sig_idx,
proc_input, success_codes
)
return asyncio.create_task(coro)
async def close(self) -> None: