scalebox-sdk 0.1.24__py3-none-any.whl → 1.0.1__py3-none-any.whl

scalebox/sandbox_async/commands/command_handle.py

@@ -1,187 +1,187 @@
-import asyncio
-import inspect
-from typing import Any, AsyncIterator, Callable, Coroutine, Optional, Tuple, Union
-
-from ...generated import api_pb2
-from ...generated.rpc import handle_rpc_exception
-from ...sandbox.commands.command_handle import (
-    CommandExitException,
-    CommandResult,
-    PtyOutput,
-    Stderr,
-    Stdout,
-)
-from ...sandbox_async.utils import OutputHandler
-
-
-class AsyncCommandHandle:
-    """
-    Command execution handle.
-
-    It provides methods for waiting for the command to finish, retrieving stdout/stderr, and killing the command.
-    """
-
-    @property
-    def pid(self):
-        """
-        Command process ID.
-        """
-        return self._pid
-
-    @property
-    def stdout(self):
-        """
-        Command stdout output.
-        """
-        return self._stdout
-
-    @property
-    def stderr(self):
-        """
-        Command stderr output.
-        """
-        return self._stderr
-
-    @property
-    def error(self):
-        """
-        Command execution error message.
-        """
-        if self._result is None:
-            return None
-        return self._result.error
-
-    @property
-    def exit_code(self):
-        """
-        Command execution exit code.
-
-        `0` if the command finished successfully.
-
-        It is `None` if the command is still running.
-        """
-        if self._result is None:
-            return None
-        return self._result.exit_code
-
-    def __init__(
-        self,
-        pid: int,
-        handle_kill: Callable[[], Coroutine[Any, Any, bool]],
-        events: AsyncIterator[
-            Union[api_pb2.StartResponse, api_pb2.ConnectResponse]
-        ],
-        on_stdout: Optional[OutputHandler[Stdout]] = None,
-        on_stderr: Optional[OutputHandler[Stderr]] = None,
-        on_pty: Optional[OutputHandler[PtyOutput]] = None,
-    ):
-        self._pid = pid
-        self._handle_kill = handle_kill
-        self._events = events
-
-        self._stdout: str = ""
-        self._stderr: str = ""
-
-        self._on_stdout = on_stdout
-        self._on_stderr = on_stderr
-        self._on_pty = on_pty
-
-        self._result: Optional[CommandResult] = None
-        self._iteration_exception: Optional[Exception] = None
-
-        self._wait = asyncio.create_task(self._handle_events())
-
-    async def _iterate_events(
-        self,
-    ) -> AsyncIterator[
-        Union[
-            Tuple[Stdout, None, None],
-            Tuple[None, Stderr, None],
-            Tuple[None, None, PtyOutput],
-        ],
-    ]:
-        async for event in self._events:
-            if event.event.HasField("data"):
-                if event.event.data.stdout:
-                    out = event.event.data.stdout.decode("utf-8", "replace")
-                    self._stdout += out
-                    yield out, None, None
-                if event.event.data.stderr:
-                    out = event.event.data.stderr.decode("utf-8", "replace")
-                    self._stderr += out
-                    yield None, out, None
-                if event.event.data.pty:
-                    yield None, None, event.event.data.pty
-            if event.event.HasField("end"):
-                self._result = CommandResult(
-                    stdout=self._stdout,
-                    stderr=self._stderr,
-                    exit_code=event.event.end.exit_code,
-                    error=event.event.end.error,
-                )
-
-    async def disconnect(self) -> None:
-        """
-        Disconnects from the command.
-
-        The command is not killed, but SDK stops receiving events from the command.
-        You can reconnect to the command using `sandbox.commands.connect` method.
-        """
-        self._wait.cancel()
-        # BUG: In Python 3.8 closing async generator can throw RuntimeError.
-        # await self._events.aclose()
-
-    async def _handle_events(self):
-        try:
-            async for stdout, stderr, pty in self._iterate_events():
-                if stdout is not None and self._on_stdout:
-                    cb = self._on_stdout(stdout)
-                    if inspect.isawaitable(cb):
-                        await cb
-                elif stderr is not None and self._on_stderr:
-                    cb = self._on_stderr(stderr)
-                    if inspect.isawaitable(cb):
-                        await cb
-                elif pty is not None and self._on_pty:
-                    cb = self._on_pty(pty)
-                    if inspect.isawaitable(cb):
-                        await cb
-        except StopAsyncIteration:
-            pass
-        except Exception as e:
-            self._iteration_exception = handle_rpc_exception(e)
-
-    async def wait(self) -> CommandResult:
-        """
-        Wait for the command to finish and return the result.
-        If the command exits with a non-zero exit code, it throws a `CommandExitException`.
-
-        :return: `CommandResult` result of command execution
-        """
-        await self._wait
-        if self._iteration_exception:
-            raise self._iteration_exception
-
-        if self._result is None:
-            raise Exception("Command ended without an end event")
-
-        if self._result.exit_code != 0:
-            raise CommandExitException(
-                stdout=self._stdout,
-                stderr=self._stderr,
-                exit_code=self._result.exit_code,
-                error=self._result.error,
-            )
-
-        return self._result
-
-    async def kill(self) -> bool:
-        """
-        Kills the command.
-
-        It uses `SIGKILL` signal to kill the command
-
-        :return: `True` if the command was killed successfully, `False` if the command was not found
-        """
-        result = await self._handle_kill()
-        return result
+import asyncio
+import inspect
+from typing import Any, AsyncIterator, Callable, Coroutine, Optional, Tuple, Union
+
+from ...generated import api_pb2
+from ...generated.rpc import handle_rpc_exception
+from ...sandbox.commands.command_handle import (
+    CommandExitException,
+    CommandResult,
+    PtyOutput,
+    Stderr,
+    Stdout,
+)
+from ...sandbox_async.utils import OutputHandler
+
+
+class AsyncCommandHandle:
+    """
+    Command execution handle.
+
+    It provides methods for waiting for the command to finish, retrieving stdout/stderr, and killing the command.
+    """
+
+    @property
+    def pid(self):
+        """
+        Command process ID.
+        """
+        return self._pid
+
+    @property
+    def stdout(self):
+        """
+        Command stdout output.
+        """
+        return self._stdout
+
+    @property
+    def stderr(self):
+        """
+        Command stderr output.
+        """
+        return self._stderr
+
+    @property
+    def error(self):
+        """
+        Command execution error message.
+        """
+        if self._result is None:
+            return None
+        return self._result.error
+
+    @property
+    def exit_code(self):
+        """
+        Command execution exit code.
+
+        `0` if the command finished successfully.
+
+        It is `None` if the command is still running.
+        """
+        if self._result is None:
+            return None
+        return self._result.exit_code
+
+    def __init__(
+        self,
+        pid: int,
+        handle_kill: Callable[[], Coroutine[Any, Any, bool]],
+        events: AsyncIterator[
+            Union[api_pb2.StartResponse, api_pb2.ConnectResponse]
+        ],
+        on_stdout: Optional[OutputHandler[Stdout]] = None,
+        on_stderr: Optional[OutputHandler[Stderr]] = None,
+        on_pty: Optional[OutputHandler[PtyOutput]] = None,
+    ):
+        self._pid = pid
+        self._handle_kill = handle_kill
+        self._events = events
+
+        self._stdout: str = ""
+        self._stderr: str = ""
+
+        self._on_stdout = on_stdout
+        self._on_stderr = on_stderr
+        self._on_pty = on_pty
+
+        self._result: Optional[CommandResult] = None
+        self._iteration_exception: Optional[Exception] = None
+
+        self._wait = asyncio.create_task(self._handle_events())
+
+    async def _iterate_events(
+        self,
+    ) -> AsyncIterator[
+        Union[
+            Tuple[Stdout, None, None],
+            Tuple[None, Stderr, None],
+            Tuple[None, None, PtyOutput],
+        ],
+    ]:
+        async for event in self._events:
+            if event.event.HasField("data"):
+                if event.event.data.stdout:
+                    out = event.event.data.stdout.decode("utf-8", "replace")
+                    self._stdout += out
+                    yield out, None, None
+                if event.event.data.stderr:
+                    out = event.event.data.stderr.decode("utf-8", "replace")
+                    self._stderr += out
+                    yield None, out, None
+                if event.event.data.pty:
+                    yield None, None, event.event.data.pty
+            if event.event.HasField("end"):
+                self._result = CommandResult(
+                    stdout=self._stdout,
+                    stderr=self._stderr,
+                    exit_code=event.event.end.exit_code,
+                    error=event.event.end.error,
+                )
+
+    async def disconnect(self) -> None:
+        """
+        Disconnects from the command.
+
+        The command is not killed, but SDK stops receiving events from the command.
+        You can reconnect to the command using `sandbox.commands.connect` method.
+        """
+        self._wait.cancel()
+        # BUG: In Python 3.8 closing async generator can throw RuntimeError.
+        # await self._events.aclose()
+
+    async def _handle_events(self):
+        try:
+            async for stdout, stderr, pty in self._iterate_events():
+                if stdout is not None and self._on_stdout:
+                    cb = self._on_stdout(stdout)
+                    if inspect.isawaitable(cb):
+                        await cb
+                elif stderr is not None and self._on_stderr:
+                    cb = self._on_stderr(stderr)
+                    if inspect.isawaitable(cb):
+                        await cb
+                elif pty is not None and self._on_pty:
+                    cb = self._on_pty(pty)
+                    if inspect.isawaitable(cb):
+                        await cb
+        except StopAsyncIteration:
+            pass
+        except Exception as e:
+            self._iteration_exception = handle_rpc_exception(e)
+
+    async def wait(self) -> CommandResult:
+        """
+        Wait for the command to finish and return the result.
+        If the command exits with a non-zero exit code, it throws a `CommandExitException`.
+
+        :return: `CommandResult` result of command execution
+        """
+        await self._wait
+        if self._iteration_exception:
+            raise self._iteration_exception
+
+        if self._result is None:
+            raise Exception("Command ended without an end event")
+
+        if self._result.exit_code != 0:
+            raise CommandExitException(
+                stdout=self._stdout,
+                stderr=self._stderr,
+                exit_code=self._result.exit_code,
+                error=self._result.error,
+            )
+
+        return self._result
+
+    async def kill(self) -> bool:
+        """
+        Kills the command.
+
+        It uses `SIGKILL` signal to kill the command
+
+        :return: `True` if the command was killed successfully, `False` if the command was not found
+        """
+        result = await self._handle_kill()
+        return result