shelltastic 0.7.0__tar.gz → 0.8.1__tar.gz

{shelltastic-0.7.0 → shelltastic-0.8.1}/PKG-INFO

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

{shelltastic-0.7.0 → shelltastic-0.8.1}/pyproject.toml

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

{shelltastic-0.7.0 → shelltastic-0.8.1}/src/shelltastic/backend/__init__.py

@@ -4,6 +4,7 @@ from shelltastic.backend.base import (
     RemoteShellBackend,
     ShellBackend,
     add_shim,
+    clear_shims,
 )
 from shelltastic.backend.default.local import DefaultLocalBackend
 from shelltastic.backend.default.remote import DefaultRemoteBackend
@@ -17,6 +18,7 @@ __all__ = [
     "register_default_local_backend",
     "register_default_remote_backend",
     "add_shim",
+    "clear_shims",
     "get_local_backend",
     "get_remote_backend",
 ]

{shelltastic-0.7.0 → shelltastic-0.8.1}/src/shelltastic/backend/base.py

@@ -90,6 +90,10 @@ class BackendShim(ABC):
 
 _backend_shims: list[BackendShim] = []
 
+def clear_shims():
+    """Clear all shims"""
+    global _backend_shims
+    _backend_shims.clear()
 
 def add_shim(shim: BackendShim):
     """Add a shim to manipulate `ShellBackend` params before any `ShellBackend` runs"""
@@ -103,7 +107,19 @@ def _shim_run(f, backend_type: type):
         result: BackendShimResult = BackendShimResult(*args, **kwargs)
         for shim in _backend_shims:
             result = shim.run_shim(
-                backend_type, f.__name__, *result.args, **result.kwargs
+                backend_type, *result.args, **result.kwargs
+            )
+        return f(*result.args, **result.kwargs)
+
+    return wrapper
+
+def _shim_run_async(f, backend_type: type):
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        result: BackendShimResult = BackendShimResult(*args, **kwargs)
+        for shim in _backend_shims:
+            result = shim.run_async_shim(
+                backend_type, *result.args, **result.kwargs
             )
         return f(*result.args, **result.kwargs)
 
@@ -114,7 +130,7 @@ class _InstallShims(ABCMeta):
     def __call__(self, *args: Any, **kwds: Any) -> Any:
         instance = super().__call__(*args, **kwds)
         instance.run = _shim_run(instance.run, instance.__class__)
-        instance.run_async = _shim_run(instance.run_async, instance.__class__)
+        instance.run_async = _shim_run_async(instance.run_async, instance.__class__)
         return instance
 
 

{shelltastic-0.7.0 → shelltastic-0.8.1}/src/shelltastic/backend/default/common.py

@@ -157,7 +157,7 @@ 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],
@@ -185,7 +185,7 @@ class CommonShellAcyncResult(ShellResultAsync[RawIOBase | None, RawIOBase | None
                 if _display_pool is None:
                     _display_pool = ThreadPool(2)
 
-                _stdout_async_result=_display_pool.apply_async(
+                _stdout_async_result = _display_pool.apply_async(
                     _output_and_pipe,
                     args=[
                         popen.stdout,
@@ -231,20 +231,16 @@ class CommonShellAcyncResult(ShellResultAsync[RawIOBase | None, RawIOBase | None
         # Wait for command to finish before closing everything
         self.wait()
 
-        print("closing popen")
         # Close popen
         self._popen.__exit__(exc_type, exc, tb)
 
         if self._display_pool:
-            print("closing _display_pool")
             self._display_pool.__exit__(exc_type, exc, tb)
 
         # Close stdout and stderr
-        if self.stdout:
-            print("closing stdout")
+        if self.stdout is not None:
             self.stdout.close()
-        if self.stderr:
-            print("closing stderr")
+        if self.stderr is not None:
             self.stderr.close()
 
     def close(self) -> None:
@@ -263,7 +259,7 @@ class CommonShellAcyncResult(ShellResultAsync[RawIOBase | None, RawIOBase | None
             stderr=self.stderr,
         )
 
-        # Wait for our async reader to be complete
+        # Wait for our async readers to be complete
         if self._stdout_async_result:
             self._stdout_async_result.wait()
 
@@ -373,7 +369,7 @@ class CommonDefaultBackend(ShellBackend):
             cwd=cwd,
         )
 
-        return CommonShellAcyncResult(
+        return CommonShellAsyncResult(
             cmd=cmd,
             popen=popen,
             stdout_capture=stdout,

shelltastic-0.8.1/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.7.0 → shelltastic-0.8.1}/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.7.0 → shelltastic-0.8.1}/src/shelltastic/result.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from io import RawIOBase
@@ -15,7 +17,15 @@ class ShellResult[SO: RawIOBase | bytes | None, SE: RawIOBase | bytes | None]:
     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)
+            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):

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

@@ -1,69 +0,0 @@
-import logging
-import shlex
-from pathlib import Path
-
-from shelltastic.backend.base import RemoteShellBackend
-from shelltastic.backend.default.common import CommonDefaultBackend
-from shelltastic.enum import CaptureMode
-from shelltastic.exception import ShellException, SSHConnectionError
-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