Fix chroot_exe when used with context manager (#156)

* Fix chroot_exe when used with context manager

Use None for `chroot_exe` default argument.  This allows us to use
similar logic as `chroot_path`.

* Update exec_helpers/api.py

Co-authored-by: Alexey Stepanov <penguinolog@users.noreply.github.com>

* Update exec_helpers/async_api/api.py

Co-authored-by: Alexey Stepanov <penguinolog@users.noreply.github.com>

---------

Co-authored-by: Geert Kloosterman <geert.kloosterman@brightcomputing.com>
Co-authored-by: Alexey Stepanov <penguinolog@users.noreply.github.com>

Author: 3 people

exec_helpers/_helpers.py

@@ -78,18 +78,19 @@ def cmd_to_string(command: str | Iterable[str]) -> str:
     return shlex.join(command)
 
 
-def chroot_command(command: str, chroot_path: str | None = None, chroot_exe: str = "chroot") -> str:
+def chroot_command(command: str, chroot_path: str | None = None, chroot_exe: str | None = None) -> str:
     """Prepare command for chroot execution.
 
     :param command: Original command.
     :type command: str
     :param chroot_path: chroot path.
     :type chroot_path: str | None
     :param chroot_exe: chroot executable.
-    :type chroot_exe: str
+    :type chroot_exe: str | None
     :return: Command to be executed with chroot rules if applicable.
     :rtype: str
     """
+    chroot_exe = chroot_exe or "chroot"
     if chroot_path and chroot_path != "/":
         chroot_dst: str = shlex.quote(chroot_path.strip())
         quoted_command = shlex.quote(command)

exec_helpers/_ssh_base.py

@@ -982,14 +982,15 @@ def keepalive(self, enforce: KeepAlivePeriodT = 1) -> _KeepAliveContext:
         """
         return _KeepAliveContext(ssh=self, enforce=int(enforce))
 
-    def _prepare_command(self, cmd: str, chroot_path: str | None = None, chroot_exe: str = "chroot") -> str:
+    def _prepare_command(self, cmd: str, chroot_path: str | None = None, chroot_exe: str | None = None) -> str:
         """Prepare command: cover chroot and other cases.
 
         :param cmd: Main command.
         :param chroot_path: Path to make chroot for execution.
         :param chroot_exe: chroot executable, default "chroot".
         :return: The final command includes chroot, if required.
         """
+        chroot_exe = chroot_exe or "chroot"
         if not self.sudo_mode:
             return super()._prepare_command(cmd=cmd, chroot_path=chroot_path, chroot_exe=chroot_exe)
         quoted_command: str = shlex.quote(cmd)
@@ -1096,7 +1097,7 @@ def open_execute_context(
         open_stdout: bool = True,
         open_stderr: bool = True,
         chroot_path: str | None = None,
-        chroot_exe: str = "chroot",
+        chroot_exe: str | None = None,
         get_pty: bool = False,
         width: int = 80,
         height: int = 24,
@@ -1116,7 +1117,7 @@ def open_execute_context(
         :param chroot_path: chroot path override.
         :type chroot_path: str | None
         :param chroot_exe: chroot exe override.
-        :type chroot_exe: str
+        :type chroot_exe: str | None
         :param get_pty: Get PTY for connection.
         :type get_pty: bool
         :param width: PTY width.
@@ -1160,7 +1161,7 @@ def execute(
         open_stderr: bool = True,
         log_stderr: bool = True,
         chroot_path: str | None = None,
-        chroot_exe: str = "chroot",
+        chroot_exe: str | None = None,
         get_pty: bool = False,
         width: int = 80,
         height: int = 24,
@@ -1190,7 +1191,7 @@ def execute(
         :param chroot_path: chroot path override.
         :type chroot_path: str | None
         :param chroot_exe: chroot exe override.
-        :type chroot_exe: str
+        :type chroot_exe: str | None
         :param get_pty: Get PTY for connection.
         :type get_pty: bool
         :param width: PTY width.
@@ -1240,7 +1241,7 @@ def __call__(
         open_stderr: bool = True,
         log_stderr: bool = True,
         chroot_path: str | None = None,
-        chroot_exe: str = "chroot",
+        chroot_exe: str | None = None,
         get_pty: bool = False,
         width: int = 80,
         height: int = 24,
@@ -1270,7 +1271,7 @@ def __call__(
         :param chroot_path: Chroot path override.
         :type chroot_path: str | None
         :param chroot_exe: Chroot exe override.
-        :type chroot_exe: str
+        :type chroot_exe: str | None
         :param get_pty: Get PTY for connection.
         :type get_pty: bool
         :param width: PTY width.
@@ -1658,7 +1659,7 @@ def execute_together(
         open_stdout: bool = True,
         open_stderr: bool = True,
         chroot_path: str | None = None,
-        chroot_exe: str = "chroot",
+        chroot_exe: str | None = None,
         verbose: bool = False,
         log_mask_re: LogMaskReT = None,
         exception_class: type[exceptions.ParallelCallProcessError] = exceptions.ParallelCallProcessError,
@@ -1685,7 +1686,7 @@ def execute_together(
         :param chroot_path: chroot path override.
         :type chroot_path: str | None
         :param chroot_exe: chroot exe override.
-        :type chroot_exe: str
+        :type chroot_exe: str | None
         :param verbose: Produce verbose log record on command call.
         :type verbose: bool
         :param log_mask_re: Regex lookup rule to mask command for logger.

exec_helpers/api.py

@@ -173,32 +173,32 @@ class _ChRootContext(typing.ContextManager[None]):
     :param path: chroot path or None for no chroot.
     :type path: str | pathlib.Path | None
     :param chroot_exe: chroot executable.
-    :type chroot_exe: str
+    :type chroot_exe: str | None
     :raises TypeError: Incorrect type of path or chroot_exe variable.
 
     .. versionadded:: 4.1.0
     """
 
     __slots__ = ("_chroot_exe_status", "_chroot_status", "_conn", "_exe", "_path")
 
-    def __init__(self, conn: ExecHelper, path: ChRootPathSetT = None, chroot_exe: str = "chroot") -> None:
+    def __init__(self, conn: ExecHelper, path: ChRootPathSetT = None, chroot_exe: str | None = None) -> None:
         """Context manager for call commands with sudo.
 
         :raises TypeError: Incorrect type of path or chroot_exe variable.
         """
         self._conn: ExecHelper = conn
         self._chroot_status: str | None = conn._chroot_path
-        self._chroot_exe_status: str = conn._chroot_exe
+        self._chroot_exe_status: str | None = conn._chroot_exe
         if path is None or isinstance(path, str):
             self._path: str | None = path
         elif isinstance(path, pathlib.Path):
             self._path = path.as_posix()  # get absolute path
         else:
             raise TypeError(f"path={path!r} is not instance of {ChRootPathSetT}")
-        if isinstance(chroot_exe, str):
-            self._exe: str = chroot_exe
+        if chroot_exe is None or isinstance(chroot_exe, str):
+            self._exe: str | None = chroot_exe
         else:
-            raise TypeError(f"chroot_exe={chroot_exe!r} is not instance of str")
+            raise TypeError(f"chroot_exe={chroot_exe!r} is not None or instance of str")
 
     def __enter__(self) -> None:
         self._conn.__enter__()
@@ -252,7 +252,7 @@ def __init__(self, log_mask_re: LogMaskReT = None, *, logger: logging.Logger) ->
         self.__logger: logging.Logger = logger
         self.log_mask_re: LogMaskReT = log_mask_re
         self.__chroot_path: str | None = None
-        self.__chroot_exe: str = "chroot"
+        self.__chroot_exe: str | None = None
         self.__context_count = 0
 
     @property
@@ -306,43 +306,43 @@ def _chroot_path(self) -> None:
         self.__chroot_path = None
 
     @property
-    def _chroot_exe(self) -> str:
+    def _chroot_exe(self) -> str | None:
         """Exe for chroot
 
-        :rtype: str
+        :rtype: str | None
         .. versionadded:: 8.1.0
         """
         return self.__chroot_exe
 
     @_chroot_exe.setter
-    def _chroot_exe(self, new_state: str) -> None:
+    def _chroot_exe(self, new_state: str | None) -> None:
         """Executable for chroot if set.
 
         :param new_state: New exe.
-        :type new_state: str
+        :type new_state: str | None
         :raises TypeError: Not supported exe information.
         .. versionadded:: 8.1.0
         """
-        if isinstance(new_state, str):
+        if new_state is None or isinstance(new_state, str):
             self.__chroot_exe = new_state
         else:
-            raise TypeError(f"chroot_exe is expected to be string, but set {new_state!r}")
+            raise TypeError(f"chroot_exe is expected to be None or string, but set {new_state!r}")
 
     @_chroot_exe.deleter
     def _chroot_exe(self) -> None:
         """Restore chroot executable.
 
         .. versionadded:: 8.1.0
         """
-        self.__chroot_exe = "chroot"
+        self.__chroot_exe = None
 
-    def chroot(self, path: ChRootPathSetT, chroot_exe: str = "chroot") -> _ChRootContext:
+    def chroot(self, path: ChRootPathSetT, chroot_exe: str | None = None) -> _ChRootContext:
         """Context manager for changing chroot rules.
 
         :param path: chroot path or none for working without chroot.
         :type path: str | pathlib.Path | None
         :param chroot_exe: chroot exe.
-        :type chroot_exe: str
+        :type chroot_exe: str | None
         :return: Context manager with selected chroot state inside.
         :rtype: typing.ContextManager
 
@@ -394,7 +394,7 @@ def _mask_command(self, cmd: str, log_mask_re: LogMaskReT = None) -> str:
 
         return _helpers.mask_command(cmd.rstrip(), self.log_mask_re, log_mask_re)
 
-    def _prepare_command(self, cmd: str, chroot_path: str | None = None, chroot_exe: str = "chroot") -> str:
+    def _prepare_command(self, cmd: str, chroot_path: str | None = None, chroot_exe: str | None = None) -> str:
         """Prepare command: cower chroot and other cases.
 
         :param cmd: Main command.
@@ -404,7 +404,11 @@ def _prepare_command(self, cmd: str, chroot_path: str | None = None, chroot_exe:
         :return: Final command, includes chroot, if required.
         :rtype: str
         """
-        return _helpers.chroot_command(cmd, chroot_path=chroot_path or self._chroot_path, chroot_exe=chroot_exe)
+        return _helpers.chroot_command(
+            cmd,
+            chroot_path=chroot_path or self._chroot_path,
+            chroot_exe=chroot_exe or self._chroot_exe,
+        )
 
     @abc.abstractmethod
     def _exec_command(
@@ -473,7 +477,7 @@ def open_execute_context(
         open_stdout: bool = True,
         open_stderr: bool = True,
         chroot_path: str | None = None,
-        chroot_exe: str = "chroot",
+        chroot_exe: str | None = None,
         **kwargs: typing.Any,
     ) -> ExecuteContext:
         """Get execution context manager.
@@ -489,7 +493,7 @@ def open_execute_context(
         :param chroot_path: chroot path override.
         :type chroot_path: str | None
         :param chroot_exe: chroot exe override.
-        :type chroot_exe: str
+        :type chroot_exe: str | None
         :param kwargs: Additional parameters for call.
         :type kwargs: typing.Any
 
@@ -509,7 +513,7 @@ def execute(
         open_stderr: bool = True,
         log_stderr: bool = True,
         chroot_path: str | None = None,
-        chroot_exe: str = "chroot",
+        chroot_exe: str | None = None,
         **kwargs: typing.Any,
     ) -> exec_result.ExecResult:
         """Execute command and wait for return code.
@@ -536,7 +540,7 @@ def execute(
         :param chroot_path: chroot path override.
         :type chroot_path: str | None
         :param chroot_exe: chroot exe override.
-        :type chroot_exe: str
+        :type chroot_exe: str | None
         :param kwargs: Additional parameters for call.
         :type kwargs: typing.Any
         :return: Execution result.
@@ -594,7 +598,7 @@ def __call__(
         open_stderr: bool = True,
         log_stderr: bool = True,
         chroot_path: str | None = None,
-        chroot_exe: str = "chroot",
+        chroot_exe: str | None = None,
         **kwargs: typing.Any,
     ) -> exec_result.ExecResult:
         """Execute command and wait for return code.
@@ -621,7 +625,7 @@ def __call__(
         :param chroot_path: chroot path override.
         :type chroot_path: str | None
         :param chroot_exe: chroot exe override.
-        :type chroot_exe: str
+        :type chroot_exe: str | None
         :param kwargs: Additional parameters for call.
         :type kwargs: typing.Any
         :return: Execution result.