feat: cache commands and remember host in exceptions

.gitignore

@@ -22,5 +22,6 @@ tests/.cache/*
 *.pot
 /*venv*
 *.mypy_cache
+.eggs
 /plumbum/version.py
 /tests/nohup.out

plumbum/commands/base.py

@@ -335,7 +335,7 @@ def popen(self, args=(), **kwargs):
 
 
 class BoundEnvCommand(BaseCommand):
-    __slots__ = ("cmd",)
+    __slots__ = ("cmd", "env", "cwd")
 
     def __init__(self, cmd, env=None, cwd=None):
         self.cmd = cmd

plumbum/commands/processes.py

@@ -1,7 +1,6 @@
 import atexit
 import heapq
 import time
-from io import StringIO
 from queue import Empty as QueueEmpty
 from queue import Queue
 from threading import Thread
@@ -114,9 +113,10 @@ class ProcessExecutionError(EnvironmentError):
     well as the command line used to create the process (``argv``)
     """
 
-    def __init__(self, argv, retcode, stdout, stderr, message=None):
+    def __init__(self, argv, retcode, stdout, stderr, message=None, host=None):
         super().__init__(self, argv, retcode, stdout, stderr)
         self.message = message
+        self.host = host
         self.argv = argv
         self.retcode = retcode
         if isinstance(stdout, bytes):
@@ -140,6 +140,8 @@ def __str__(self):
             lines = ["Unexpected exit code: ", str(self.retcode)]
         cmd = "\n              | ".join(cmd.splitlines())
         lines += ["\nCommand line: | ", cmd]
+        if self.host:
+            lines += ["\nHost:         | ", self.host]
         if stdout:
             lines += ["\nStdout:       | ", stdout]
         if stderr:
@@ -304,6 +306,7 @@ def run_proc(proc, retcode, timeout=None):
 BY_POSITION = object()
 BY_TYPE = object()
 DEFAULT_ITER_LINES_MODE = BY_POSITION
+DEFAULT_BUFFER_SIZE = _INFINITE = float("inf")
 
 
 def iter_lines(
@@ -312,6 +315,7 @@ def iter_lines(
     timeout=None,
     linesize=-1,
     line_timeout=None,
+    buffer_size=None,
     mode=None,
     _iter_lines=_iter_lines,
 ):
@@ -335,25 +339,41 @@ def iter_lines(
                     Raise an :class:`ProcessLineTimedOut <plumbum.commands.ProcessLineTimedOut>` if the timeout has
                     been reached. ``None`` means no timeout is imposed.
 
+    :param buffer_size: Maximum number of lines to keep in the stdout/stderr buffers, in case of a ProcessExecutionError.
+                    Default is ``None``, which defaults to DEFAULT_BUFFER_SIZE (which is infinite by default).
+                    ``0`` will disable bufferring completely.
+
+    :param mode: Controls what the generator yields. Defaults to DEFAULT_ITER_LINES_MODE (which is BY_POSITION by default)
+                - BY_POSITION (default): yields ``(out, err)`` line tuples, where either item may be ``None``
+                - BY_TYPE: yields ``(fd, line)`` tuples, where ``fd`` is 1 (stdout) or 2 (stderr)
+
     :returns: An iterator of (out, err) line tuples.
     """
     if mode is None:
         mode = DEFAULT_ITER_LINES_MODE
 
+    if buffer_size is None:
+        buffer_size = DEFAULT_BUFFER_SIZE
+    buffer_size: int
+
     assert mode in (BY_POSITION, BY_TYPE)
 
     encoding = getattr(proc, "custom_encoding", None) or "utf-8"
     decode = lambda s: s.decode(encoding, errors="replace").rstrip()
 
     _register_proc_timeout(proc, timeout)
 
-    buffers = [StringIO(), StringIO()]
+    buffers = [[], []]
     for t, line in _iter_lines(proc, decode, linesize, line_timeout):
 
         # verify that the proc hasn't timed out yet
         proc.verify(timeout=timeout, retcode=None, stdout=None, stderr=None)
 
-        buffers[t].write(line + "\n")
+        buffer = buffers[t]
+        if buffer_size > 0:
+            buffer.append(line)
+            if buffer_size < _INFINITE:
+                del buffer[:-buffer_size]
 
         if mode is BY_POSITION:
             ret = [None, None]
@@ -363,4 +383,4 @@ def iter_lines(
             yield (t + 1), line  # 1=stdout, 2=stderr
 
     # this will take care of checking return code and timeouts
-    _check_process(proc, retcode, timeout, *(s.getvalue() for s in buffers))
+    _check_process(proc, retcode, timeout, *("\n".join(s) + "\n" for s in buffers))

plumbum/machines/base.py

@@ -91,3 +91,12 @@ def __getattr__(self, name):
     @property
     def cmd(self):
         return self.Cmd(self)
+
+    def clear_program_cache(self):
+        """
+        Clear the program cache, which is populated via ``machine.which(progname)`` calls.
+
+        This cache speeds up the lookup of a program in the machines PATH, and is particularly
+        effective for RemoteMachines.
+        """
+        self._program_cache.clear()

plumbum/machines/local.py

@@ -8,6 +8,7 @@
 from contextlib import contextmanager
 from subprocess import PIPE, Popen
 from tempfile import mkdtemp
+from typing import Dict, Tuple
 
 from plumbum.commands import CommandNotFound, ConcreteCommand
 from plumbum.commands.daemons import posix_daemonize, win32_daemonize
@@ -141,6 +142,7 @@ class LocalMachine(BaseMachine):
 
     custom_encoding = sys.getfilesystemencoding()
     uname = platform.uname()[0]
+    _program_cache: Dict[Tuple[str, str], LocalPath] = {}
 
     def __init__(self):
         self._as_user_stack = []
@@ -182,13 +184,22 @@ def which(cls, progname):
 
         :returns: A :class:`LocalPath <plumbum.machines.local.LocalPath>`
         """
+
+        key = (progname, cls.env.get("PATH", ""))
+
+        try:
+            return cls._program_cache[key]
+        except KeyError:
+            pass
+
         alternatives = [progname]
         if "_" in progname:
             alternatives.append(progname.replace("_", "-"))
             alternatives.append(progname.replace("_", "."))
         for pn in alternatives:
             path = cls._which(pn)
             if path:
+                cls._program_cache[key] = path
                 return path
         raise CommandNotFound(progname, list(cls.env.path))
 

plumbum/machines/remote.py

@@ -173,6 +173,7 @@ def __init__(self, encoding="utf8", connect_timeout=10, new_session=False):
         self.uname = self._get_uname()
         self.env = RemoteEnv(self)
         self._python = None
+        self._program_cache = {}
 
     def _get_uname(self):
         rc, out, _ = self._session.run("uname", retcode=None)
@@ -225,6 +226,13 @@ def which(self, progname):
 
         :returns: A :class:`RemotePath <plumbum.path.local.RemotePath>`
         """
+        key = (progname, self.env.get("PATH", ""))
+
+        try:
+            return self._program_cache[key]
+        except KeyError:
+            pass
+
         alternatives = [progname]
         if "_" in progname:
             alternatives.append(progname.replace("_", "-"))
@@ -233,6 +241,7 @@ def which(self, progname):
             for p in self.env.path:
                 fn = p / name
                 if fn.access("x") and not fn.is_dir():
+                    self._program_cache[key] = fn
                     return fn
 
         raise CommandNotFound(progname, self.env.path)

plumbum/machines/session.py

@@ -72,7 +72,8 @@ class SessionPopen(PopenAddons):
     """A shell-session-based ``Popen``-like object (has the following attributes: ``stdin``,
     ``stdout``, ``stderr``, ``returncode``)"""
 
-    def __init__(self, proc, argv, isatty, stdin, stdout, stderr, encoding):
+    def __init__(self, host, proc, argv, isatty, stdin, stdout, stderr, encoding):
+        self.host = host
         self.proc = proc
         self.argv = argv
         self.isatty = isatty
@@ -132,6 +133,7 @@ def communicate(self, input=None):  # pylint: disable=redefined-builtin
                         stdout,
                         stderr,
                         message="Incorrect username or password provided",
+                        host=self.host,
                     ) from None
                 if returncode == 6:
                     raise HostPublicKeyUnknown(
@@ -140,6 +142,7 @@ def communicate(self, input=None):  # pylint: disable=redefined-builtin
                         stdout,
                         stderr,
                         message="The authenticity of the host can't be established",
+                        host=self.host,
                     ) from None
                 if returncode != 0:
                     raise SSHCommsError(
@@ -148,6 +151,7 @@ def communicate(self, input=None):  # pylint: disable=redefined-builtin
                         stdout,
                         stderr,
                         message="SSH communication failed",
+                        host=self.host,
                     ) from None
                 if name == "2":
                     raise SSHCommsChannel2Error(
@@ -156,6 +160,7 @@ def communicate(self, input=None):  # pylint: disable=redefined-builtin
                         stdout,
                         stderr,
                         message="No stderr result detected. Does the remote have Bash as the default shell?",
+                        host=self.host,
                     ) from None
 
                 raise SSHCommsError(
@@ -164,6 +169,7 @@ def communicate(self, input=None):  # pylint: disable=redefined-builtin
                     stdout,
                     stderr,
                     message="No communication channel detected. Does the remote exist?",
+                    host=self.host,
                 ) from err
             if not line:
                 del sources[i]
@@ -202,7 +208,10 @@ class ShellSession:
                             is seen, the shell process is killed
     """
 
-    def __init__(self, proc, encoding="auto", isatty=False, connect_timeout=5):
+    def __init__(
+        self, proc, encoding="auto", isatty=False, connect_timeout=5, host=None
+    ):
+        self.host = host
         self.proc = proc
         self.custom_encoding = proc.custom_encoding if encoding == "auto" else encoding
         self.isatty = isatty
@@ -292,6 +301,7 @@ def popen(self, cmd):
         self.proc.stdin.write(full_cmd + b"\n")
         self.proc.stdin.flush()
         self._current = SessionPopen(
+            self.host,
             self.proc,
             full_cmd,
             self.isatty,

plumbum/machines/ssh_machine.py

@@ -104,6 +104,7 @@ def __init__(
 
         scp_args = []
         ssh_args = []
+        self.host = host
         if user:
             self._fqhost = f"{user}@{host}"
         else:
@@ -208,6 +209,7 @@ def session(self, isatty=False, new_session=False):
             self.custom_encoding,
             isatty,
             self.connect_timeout,
+            host=self.host,
         )
 
     def tunnel(

tests/test_local.py

@@ -551,6 +551,16 @@ def test_iter_lines_timeout(self):
                 print(i, "out:", out)
         assert i in (2, 3)  # Mac is a bit flakey
 
+    @skip_on_windows
+    def test_iter_lines_buffer_size(self):
+        from plumbum.cmd import bash
+
+        cmd = bash["-ce", "for ((i=0;i<100;i++)); do echo $i; done; false"]
+        with pytest.raises(ProcessExecutionError) as e:
+            for _ in cmd.popen().iter_lines(timeout=1, buffer_size=5):
+                pass
+        assert e.value.stdout == "\n".join(map(str, range(95, 100))) + "\n"
+
     @skip_on_windows
     def test_iter_lines_timeout_by_type(self):
         from plumbum.cmd import bash