diff --git a/.gitignore b/.gitignore index 3905ee6b..6870fe28 100644 --- a/.gitignore +++ b/.gitignore @@ -22,5 +22,6 @@ tests/.cache/* *.pot /*venv* *.mypy_cache +.eggs /plumbum/version.py /tests/nohup.out diff --git a/plumbum/commands/processes.py b/plumbum/commands/processes.py index c0901017..da33b76e 100644 --- a/plumbum/commands/processes.py +++ b/plumbum/commands/processes.py @@ -113,13 +113,14 @@ class ProcessExecutionError(OSError): 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): # we can't use 'super' here since OSError only keeps the first 2 args, # which leads to failuring in loading this object from a pickle.dumps. Exception.__init__(self, argv, retcode, stdout, stderr) self.message = message + self.host = host self.argv = argv self.retcode = retcode if isinstance(stdout, bytes): @@ -143,6 +144,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: diff --git a/plumbum/machines/base.py b/plumbum/machines/base.py index a498c7c2..754852be 100644 --- a/plumbum/machines/base.py +++ b/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() diff --git a/plumbum/machines/local.py b/plumbum/machines/local.py index 723c72c5..a9a4e5f4 100644 --- a/plumbum/machines/local.py +++ b/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,6 +184,14 @@ def which(cls, progname): :returns: A :class:`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("_", "-")) @@ -189,6 +199,7 @@ def which(cls, progname): for pn in alternatives: path = cls._which(pn) if path: + cls._program_cache[key] = path return path raise CommandNotFound(progname, list(cls.env.path)) diff --git a/plumbum/machines/remote.py b/plumbum/machines/remote.py index df86fa42..097652d5 100644 --- a/plumbum/machines/remote.py +++ b/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 ` """ + 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) diff --git a/plumbum/machines/session.py b/plumbum/machines/session.py index 21d13217..f292eded 100644 --- a/plumbum/machines/session.py +++ b/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, proc, argv, isatty, stdin, stdout, stderr, encoding, *, host): + 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 @@ -299,6 +308,7 @@ def popen(self, cmd): MarkedPipe(self.proc.stdout, marker), MarkedPipe(self.proc.stderr, marker), self.custom_encoding, + host=self.host, ) return self._current diff --git a/plumbum/machines/ssh_machine.py b/plumbum/machines/ssh_machine.py index 8ce41c43..648b0efb 100644 --- a/plumbum/machines/ssh_machine.py +++ b/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(