shelltastic 0.6.0__tar.gz → 0.8.0__tar.gz

{shelltastic-0.6.0 → shelltastic-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shelltastic
-Version: 0.6.0
+Version: 0.8.0
 Summary: A fantastic shell command runner for python
 Author: Bearmine
 License-Expression: MPL-2.0

{shelltastic-0.6.0 → shelltastic-0.8.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "shelltastic"
-version = "0.6.0"
+version = "0.8.0"
 description = "A fantastic shell command runner for python"
 readme = "README.md"
 authors = [

{shelltastic-0.6.0 → shelltastic-0.8.0}/src/shelltastic/backend/default/common.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import logging
 import os
 import shlex
@@ -7,8 +5,10 @@ import subprocess
 import time
 from io import RawIOBase
 from multiprocessing import Pipe
+from multiprocessing.connection import Connection
 from multiprocessing.pool import ThreadPool
-from typing import IO, TYPE_CHECKING, Literal
+from pathlib import Path
+from typing import IO, Literal
 
 from shelltastic import display
 from shelltastic.backend.base import ShellBackend
@@ -16,10 +16,6 @@ from shelltastic.enum import CaptureMode, DisplayMode
 from shelltastic.io import PipeReader
 from shelltastic.result import ShellResult, ShellResultAsync
 
-if TYPE_CHECKING:
-    from multiprocessing.connection import Connection
-    from pathlib import Path
-
 LOGGER = logging.getLogger(__name__)
 
 
@@ -91,45 +87,27 @@ class IOTimeoutStream:
                 return return_val
 
 
-def _output_and_collect(
-    io: IO[bytes] | None, display: display.IODisplay | None, collect: bool = True
-) -> bytes | None:
-    if io is None:
-        return None
-
-    if collect:
-        if display:
-            return_bytes: bytes = b""
-            for line in IOTimeoutStream(io):
-                return_bytes += line
-                display.printbytes(line)
-            return return_bytes
-        return io.read()
-
-    if display:
-        for line in IOTimeoutStream(io):
-            display.printbytes(line)
-
-    return None
-
-
 def _output_and_pipe(
     io: IO[bytes] | None, display: display.IODisplay | None, pipe: Connection | None
 ) -> None:
-    if io is None:
-        return
+    try:
+        if io is None:
+            return
+
+        if pipe:
+            if display:
+                for line in IOTimeoutStream(io):
+                    pipe.send_bytes(line)
+                    display.printbytes(line)
+            for line in IOTimeoutStream(io):
+                pipe.send_bytes(line)
 
-    if pipe:
         if display:
             for line in IOTimeoutStream(io):
-                pipe.send_bytes(line)
                 display.printbytes(line)
-        for line in IOTimeoutStream(io):
-            pipe.send_bytes(line)
-
-    if display:
-        for line in IOTimeoutStream(io):
-            display.printbytes(line)
+    finally:
+        if pipe:
+            pipe.close()
 
 
 def _determine_subprocess_output_mode(
@@ -179,55 +157,121 @@ def _determine_subprocess_output_mode(
     return subprocess.PIPE
 
 
-class CommonShellAcyncResult(ShellResultAsync[RawIOBase | None, RawIOBase | None]):
+class CommonShellAsyncResult(ShellResultAsync[RawIOBase | None, RawIOBase | None]):
     def __init__(
         self,
         cmd: str | list[str],
         popen: subprocess.Popen,
+        stdout_capture: CaptureMode | None,
+        stderr_capture: CaptureMode | None,
         stdout_display: display.IODisplay | None,
         stderr_display: display.IODisplay | None,
+        check: bool,
     ) -> None:
-        stdout_rcv, stdout_snd = Pipe(False)
-        stderr_rcv, stderr_snd = Pipe(False)
-        super().__init__(cmd, PipeReader(stdout_rcv), PipeReader(stderr_rcv))
-        self._popen = popen
-        self._display_pool = ThreadPool(2)
-
-        # Print out and capture stdout
-        self._stdout_result = self._display_pool.apply_async(
-            _output_and_pipe,
-            args=[
-                popen.stdout,
-                stdout_display,
-                stdout_snd,
-            ],
-        )
+        _display_pool = None
+        _stdout_async_result = None
+        _stderr_async_result = None
+
+        if popen.stdout:
+            # If stdout is PIPE, then setup reader
+            if stdout_capture == CaptureMode.PIPE:
+                stdout_rcv, stdout_snd = Pipe(False)
+                stdout = PipeReader(stdout_rcv)
+            else:
+                stdout_snd = None
+                stdout = None
+
+            if stdout_snd or stdout_display:
+                if _display_pool is None:
+                    _display_pool = ThreadPool(2)
+
+                _stdout_async_result = _display_pool.apply_async(
+                    _output_and_pipe,
+                    args=[
+                        popen.stdout,
+                        stdout_display,
+                        stdout_snd,
+                    ],
+                )
+        else:
+            stdout = None
 
-        # Print out and capture stderr
-        self._stderr_result = self._display_pool.apply_async(
-            _output_and_pipe,
-            args=[
-                popen.stderr,
-                stderr_display,
-                stderr_snd,
-            ],
-        )
+        if popen.stderr:
+            if stderr_capture == CaptureMode.PIPE:
+                stderr_rcv, stderr_snd = Pipe(False)
+                stderr = PipeReader(stderr_rcv)
+            else:
+                stderr_snd = None
+                stderr = None
+
+            if stderr_snd or stderr_display:
+                if _display_pool is None:
+                    _display_pool = ThreadPool(2)
+
+                _stderr_async_result = _display_pool.apply_async(
+                    _output_and_pipe,
+                    args=[
+                        popen.stderr,
+                        stderr_display,
+                        stderr_snd,
+                    ],
+                )
+        else:
+            stderr = None
+
+        super().__init__(cmd, stdout, stderr)
+
+        self._display_pool = _display_pool
+        self._stdout_async_result = _stdout_async_result
+        self._stderr_async_result = _stderr_async_result
+        self._popen = popen
+        self._check = check
 
     def __exit__(self, exc_type, exc, tb):
+        # Wait for command to finish before closing everything
+        self.wait()
+
+        # Close popen
         self._popen.__exit__(exc_type, exc, tb)
-        self._display_pool.__exit__(exc_type, exc, tb)
+
+        if self._display_pool:
+            self._display_pool.__exit__(exc_type, exc, tb)
+
+        # Close stdout and stderr
+        if self.stdout is not None:
+            self.stdout.close()
+        if self.stderr is not None:
+            self.stderr.close()
 
     def close(self) -> None:
         self.__exit__(None, None, None)
 
+    def terminate(self) -> None:
+        self._popen.terminate()
+        self.close()
+
     def result(self) -> ShellResult[RawIOBase | None, RawIOBase | None]:
-        return ShellResult(
+        # Get result
+        result = ShellResult(
             cmd=self.cmd,
             returncode=self._popen.wait(),
             stdout=self.stdout,
             stderr=self.stderr,
         )
 
+        # Wait for our async readers to be complete
+        if self._stdout_async_result:
+            self._stdout_async_result.wait()
+
+        if self._stderr_async_result:
+            self._stderr_async_result.wait()
+
+        # Check if set
+        if self._check:
+            result.check_returncode()
+
+        return result
+
 
 class CommonDefaultBackend(ShellBackend):
     __slots__ = ()
@@ -245,106 +289,18 @@ class CommonDefaultBackend(ShellBackend):
         echo_cmd: DisplayMode | display.IODisplay | bool | None = None,
         **kwargs,
     ) -> ShellResult[bytes | None, bytes | None]:
-        if kwargs:
-            raise TypeError(f"Uknown argument/s {list(kwargs.keys())}")
-
-        # If cmd is a str, run as shell
-        shell = isinstance(cmd, str)
-
-        # Echo command based on setting
-        # Defaults to DEBUG_LOG
-        # False to disable
-        if isinstance(echo_cmd, bool):
-            echo_cmd = DisplayMode.STDOUT if echo_cmd else DisplayMode.DEVNULL
-
-        if echo_cmd is None:
-            echo_cmd = DisplayMode.DEBUG_LOG
-
-        if isinstance(echo_cmd, display.IODisplay):
-            echo_display = echo_cmd
-        else:
-            echo_display = display.from_mode(echo_cmd)
-
-        if isinstance(cmd, str):
-            echo_display.printline("%s> %s", cwd or "", cmd)
-        else:
-            echo_display.printline("%s> %s", cwd or "", shlex.join(cmd))
-
-        # Determine PIPE mode for stdout
-        stdout_sub_mode = _determine_subprocess_output_mode(
-            stdout, stdout_display, "stdout"
-        )
-
-        # Determine PIPE mode for stderr
-        stderr_sub_mode = _determine_subprocess_output_mode(
-            stderr, stderr_display, "stderr"
-        )
-
-        if stdout_display is None:
-            stdout_display = (
-                DisplayMode.STDOUT if stdout is None else DisplayMode.DEVNULL
-            )
-
-        if stderr_display is None:
-            stderr_display = (
-                DisplayMode.STDERR if stderr is None else DisplayMode.DEVNULL
-            )
-
-        # Run subprocess command
-        with (
-            subprocess.Popen(
-                cmd,
-                shell=shell,
-                stdout=stdout_sub_mode,
-                stderr=stderr_sub_mode,
-                cwd=cwd,
-            ) as popen,
-            ThreadPool(2) as display_pool,
-        ):
-            # Print out and capture stdout
-            stdout_result = display_pool.apply_async(
-                _output_and_collect,
-                args=[
-                    popen.stdout,
-                    (
-                        stdout_display
-                        if isinstance(stdout_display, display.IODisplay)
-                        else display.from_mode(stdout_display)
-                    ),
-                    stdout is not None and stdout != CaptureMode.DEVNULL,
-                ],
-            )
-
-            # Print out and capture stderr
-            stderr_result = display_pool.apply_async(
-                _output_and_collect,
-                args=[
-                    popen.stderr,
-                    (
-                        stderr_display
-                        if isinstance(stderr_display, display.IODisplay)
-                        else display.from_mode(stderr_display)
-                    ),
-                    stderr is not None and stderr != CaptureMode.DEVNULL,
-                ],
-            )
-
-            # Wait on command completion
-            returncode = popen.wait()
-
-            # Build our result
-            result = ShellResult(
-                cmd=cmd,
-                stdout=stdout_result.get(),
-                stderr=stderr_result.get(),
-                returncode=returncode,
-            )
-
-        # If check is True, raise exception if returncode is not 0
-        if check:
-            result.check_returncode()
-
-        return result
+        with self.run_async(
+            cmd,
+            check=check,
+            stdout=stdout,
+            stderr=stderr,
+            stdout_display=stdout_display,
+            stderr_display=stderr_display,
+            cwd=cwd,
+            echo_cmd=echo_cmd,
+            **kwargs,
+        ) as running_cmd:
+            return running_cmd.result_bytes()
 
     def run_async(
         self,
@@ -413,17 +369,20 @@ class CommonDefaultBackend(ShellBackend):
             cwd=cwd,
         )
 
-        return CommonShellAcyncResult(
-            cmd,
-            popen,
-            (
+        return CommonShellAsyncResult(
+            cmd=cmd,
+            popen=popen,
+            stdout_capture=stdout,
+            stderr_capture=stderr,
+            stdout_display=(
                 stdout_display
                 if isinstance(stdout_display, display.IODisplay)
                 else display.from_mode(stdout_display)
             ),
-            (
+            stderr_display=(
                 stderr_display
                 if isinstance(stderr_display, display.IODisplay)
                 else display.from_mode(stderr_display)
             ),
+            check=check,
         )

{shelltastic-0.6.0 → shelltastic-0.8.0}/src/shelltastic/backend/default/local.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import logging
 
 from shelltastic.backend.base import LocalShellBackend

shelltastic-0.8.0/src/shelltastic/backend/default/remote.py

@@ -0,0 +1,147 @@
+import logging
+import shlex
+from io import RawIOBase
+from pathlib import Path
+from subprocess import Popen
+
+from shelltastic.backend.base import RemoteShellBackend
+from shelltastic.backend.default.common import (
+    CommonDefaultBackend,
+    CommonShellAsyncResult,
+)
+from shelltastic.display import IODisplay
+from shelltastic.enum import CaptureMode
+from shelltastic.exception import ShellException, SSHConnectionError
+from shelltastic.host import Host
+from shelltastic.result import ShellResult, ShellResultAsync
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _monkey_patch_async_result(
+    async_result: ShellResultAsync[RawIOBase | None, RawIOBase | None], host: Host
+):
+    result_fn = async_result.result
+
+    def patched_result_fn():
+        # Extra check for SSH Errors
+        try:
+            result = result_fn()
+            if (
+                result.returncode == 255
+            ):  # 255 is what ssh returns if it has an error with ssh itself
+                result = result.to_result_bytes()
+                raise SSHConnectionError(
+                    host.hostname, result.returncode, result.stdout, result.stderr
+                )
+            return result
+        except ShellException as ex:
+            if ex.returncode == 255:
+                raise SSHConnectionError(
+                    host.hostname, ex.returncode, ex.stdout, ex.stderr
+                ) from ex
+            raise
+
+    async_result.result = patched_result_fn
+
+    return async_result
+
+
+class RemoteShellAcyncResult(CommonShellAsyncResult):
+    def __init__(
+        self,
+        cmd: str | list[str],
+        popen: Popen,
+        stdout_capture: CaptureMode | None,
+        stderr_capture: CaptureMode | None,
+        stdout_display: IODisplay | None,
+        stderr_display: IODisplay | None,
+        check: bool,
+        host: Host,
+    ) -> None:
+        super().__init__(
+            cmd,
+            popen,
+            stdout_capture,
+            stderr_capture,
+            stdout_display,
+            stderr_display,
+            check,
+        )
+        self.host = host
+
+    def result(self) -> ShellResult[RawIOBase | None, RawIOBase | None]:
+        # Extra check for SSH Errors
+        try:
+            result = super().result()
+            if (
+                result.returncode == 255
+            ):  # 255 is what ssh returns if it has an error with ssh itself
+                result = result.to_result_bytes()
+                raise SSHConnectionError(
+                    self.host.hostname, result.returncode, result.stdout, result.stderr
+                )
+            return result
+        except ShellException as ex:
+            if ex.returncode == 255:
+                raise SSHConnectionError(
+                    self.host.hostname, ex.returncode, ex.stdout, ex.stderr
+                ) from ex
+            raise
+
+
+class DefaultRemoteBackend(CommonDefaultBackend, RemoteShellBackend):
+    def __init__(self, host: Host) -> None:
+        super().__init__()
+        self.host: Host = host
+
+    @staticmethod
+    def for_host(host: Host) -> RemoteShellBackend:
+        return DefaultRemoteBackend(host)
+
+    def run(
+        self,
+        cmd: str | list[str],
+        *,
+        subshell: bool = True,
+        cwd: str | Path | None = None,
+        **kwargs,
+    ) -> ShellResult:
+        """
+        Execute a remote command using SSH.
+        """
+        with self.run_async(cmd, subshell=subshell, cwd=cwd, **kwargs) as running_cmd:
+            return running_cmd.result_bytes()
+
+    def run_async(
+        self,
+        cmd: str | list[str],
+        *,
+        subshell: bool = True,
+        cwd: str | Path | None = None,
+        **kwargs,
+    ) -> ShellResultAsync[RawIOBase | None, RawIOBase | None]:
+        if isinstance(cmd, list):
+            cmd = shlex.join(cmd)
+
+        # Wrap user command in sh -c call
+        if subshell:
+            cmd = shlex.join(["bash", "-c", cmd])
+
+        # If we're redirecting stderr to stdout,
+        # add the same redirect inside the ssh command.
+        if kwargs.get("stderr") == CaptureMode.STDOUT:
+            cmd = cmd + " 2>&1"
+
+        # cd into cwd first if cwd is set
+        if cwd:
+            cmd = shlex.join(["cd", str(cwd)]) + " && " + cmd
+
+        port_flag = ["-p", str(self.host.port)] if self.host.port else []
+
+        return _monkey_patch_async_result(
+            super().run_async(
+                ["ssh", *port_flag, self.host.host_specifier(), cmd], **kwargs
+            ),
+            self.host,
+        )

{shelltastic-0.6.0 → shelltastic-0.8.0}/src/shelltastic/exception.py

@@ -1,6 +1,3 @@
-from io import RawIOBase
-
-
 class SSHConnectionError(Exception):
     def __init__(
         self, hostname: str, returncode: int, stdout: bytes | None, stderr: bytes | None
@@ -15,13 +12,13 @@ class SSHConnectionError(Exception):
         self.stderr = stderr
 
 
-class ShellException[SO: RawIOBase | bytes | None, SE: RawIOBase | bytes | None](Exception):
+class ShellException(Exception):
     def __init__(
         self,
         cmd: str | list[str],
         returncode: int,
-        stdout: SO,
-        stderr: SE,
+        stdout: bytes | None,
+        stderr: bytes | None,
     ):
         super().__init__(f"Shell execution failed with returncode {returncode}")
         self.cmd = cmd

shelltastic-0.8.0/src/shelltastic/result.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from io import RawIOBase
+
+from shelltastic.exception import ShellException
+
+
+@dataclass
+class ShellResult[SO: RawIOBase | bytes | None, SE: RawIOBase | bytes | None]:
+    cmd: str | list[str]
+    returncode: int
+    stdout: SO
+    stderr: SE
+
+    def check_returncode(self):
+        """If returncode is not 0, raise ShellException"""
+        if self.returncode != 0:
+            result = self.to_result_bytes()
+            raise ShellException(self.cmd, self.returncode, result.stdout, result.stderr)
+
+    def to_result_bytes(self) -> ShellResult[bytes | None, bytes | None]:
+        stdout = self.stdout.readall() if isinstance(self.stdout, RawIOBase) else self.stdout
+
+        stderr = self.stderr.readall() if isinstance(self.stderr, RawIOBase) else self.stderr
+
+        return ShellResult(self.cmd, self.returncode, stdout, stderr)
+
+
+class ShellResultAsync[SO: RawIOBase | None, SE: RawIOBase | None](ABC):
+    """
+    An asynchronsly running shell command.
+
+    cmd: String for the running shell command.
+
+    stdout: Bytes stream for stdout
+
+    stderr: Bytes stream for stderr
+    """
+
+    def __init__(self, cmd: str | list[str], stdout: SO, stderr: SE) -> None:
+        self.cmd = cmd
+        """The running command"""
+
+        self.stdout = stdout
+        """Byte based IO stream for stdout"""
+
+        self.stderr = stderr
+        """Byte based IO stream for stderr"""
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.close()
+
+    def close(self) -> None:
+        """Close the running command."""
+        return
+
+    def terminate(self) -> None:
+        """Terminate the running command"""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def result(self) -> ShellResult[SO, SE]:
+        """
+        Blocking
+
+        Wait for the command to finish and return a ShellResult with the outcome.
+        """
+        raise NotImplementedError()
+
+    def wait(self) -> None:
+        """
+        Blocking
+
+        Wait for the command to finish.
+        """
+        self.result()
+
+    def returncode(self) -> int:
+        """
+        Blocking
+
+        Get the returncode of the command.
+        """
+        return self.result().returncode
+
+    def check_returncode(self):
+        """
+        Blocking
+
+        Raise a ShellException if the returncode is not zero.
+        """
+        self.result().check_returncode()
+
+    def result_bytes(self) -> ShellResult[bytes | None, bytes | None]:
+        result = self.result()
+        stdout = result.stdout
+        stdout_bytes = stdout.readall() if stdout is not None else None
+        stderr = result.stderr
+        stderr_bytes = stderr.readall() if stderr is not None else None
+
+        return ShellResult(
+            cmd=result.cmd,
+            returncode=result.returncode,
+            stdout=stdout_bytes,
+            stderr=stderr_bytes,
+        )

shelltastic-0.6.0/src/shelltastic/backend/default/remote.py

@@ -1,75 +0,0 @@
-from __future__ import annotations
-
-import logging
-import shlex
-from typing import TYPE_CHECKING
-
-from shelltastic.backend.base import RemoteShellBackend
-from shelltastic.backend.default.common import CommonDefaultBackend
-from shelltastic.enum import CaptureMode
-from shelltastic.exception import ShellException, SSHConnectionError
-
-if TYPE_CHECKING:
-    from pathlib import Path
-
-    from shelltastic.host import Host
-    from shelltastic.result import ShellResult
-
-LOGGER = logging.getLogger(__name__)
-
-
-class DefaultRemoteBackend(CommonDefaultBackend, RemoteShellBackend):
-    def __init__(self, host: Host) -> None:
-        super().__init__()
-        self.host: Host = host
-
-    @staticmethod
-    def for_host(host: Host) -> RemoteShellBackend:
-        return DefaultRemoteBackend(host)
-
-    def run(
-        self,
-        cmd: str | list[str],
-        *,
-        subshell: bool = True,
-        cwd: str | Path | None = None,
-        **kwargs,
-    ) -> ShellResult:
-        """
-        Execute a remote command using SSH.
-        """
-        if isinstance(cmd, list):
-            cmd = shlex.join(cmd)
-
-        # Wrap user command in sh -c call
-        if subshell:
-            cmd = shlex.join(["bash", "-c", cmd])
-
-        # If we're redirecting stderr to stdout,
-        # add the same redirect inside the ssh command.
-        if kwargs.get("stderr") == CaptureMode.STDOUT:
-            cmd = cmd + " 2>&1"
-
-        # cd into cwd first if cwd is set
-        if cwd:
-            cmd = shlex.join(["cd", str(cwd)]) + " && " + cmd
-
-        port_flag = ["-p", str(self.host.port)] if self.host.port else []
-
-        try:
-            result = super().run(
-                ["ssh", *port_flag, self.host.host_specifier(), cmd], **kwargs
-            )
-            if (
-                result.returncode == 255
-            ):  # 255 is what ssh returns if it has an error with ssh itself
-                raise SSHConnectionError(
-                    self.host.hostname, result.returncode, result.stdout, result.stderr
-                )
-            return result
-        except ShellException as ex:
-            if ex.returncode == 255:
-                raise SSHConnectionError(
-                    self.host.hostname, ex.returncode, ex.stdout, ex.stderr
-                ) from ex
-            raise

shelltastic-0.6.0/src/shelltastic/result.py

@@ -1,47 +0,0 @@
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from io import RawIOBase
-
-from shelltastic.exception import ShellException
-
-
-@dataclass
-class ShellResult[SO: RawIOBase | bytes | None, SE: RawIOBase | bytes | None]:
-    cmd: str | list[str]
-    returncode: int
-    stdout: SO
-    stderr: SE
-
-    def check_returncode(self):
-        """If returncode is not 0, raise ShellException"""
-        if self.returncode != 0:
-            raise ShellException(self.cmd, self.returncode, self.stdout, self.stderr)
-
-
-class ShellResultAsync[SO: RawIOBase | None, SE: RawIOBase | None](ABC):
-    def __init__(self, cmd: str | list[str], stdout: SO, stderr: SE) -> None:
-        self.cmd = cmd
-        self.stdout = stdout
-        self.stderr = stderr
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc, tb):
-        self.close()
-
-    def close(self) -> None:
-        return
-
-    @abstractmethod
-    def result(self) -> ShellResult[SO, SE]:
-        raise NotImplementedError()
-
-    def wait(self) -> None:
-        self.result()
-
-    def returncode(self) -> int:
-        return self.result().returncode
-
-    def check_returncode(self):
-        self.result().check_returncode()