added output truncate to cli

bale/interfaces/cli.py

@@ -34,6 +34,7 @@ class Cli:
         self.stderr: List[str] = []
         self._terminate: asyncio.Event = asyncio.Event()
         self._busy: bool = False
+        self._truncated: bool = False
         self.prefix_line: str = ""
         self._stdout_terminals: List[Terminal] = []
         self._stderr_terminals: List[Terminal] = []
@@ -70,8 +71,11 @@ class Cli:
             else:
                 break
 
-    async def _controller(self, process: Process) -> None:
+    async def _controller(self, process: Process, max_output_lines) -> None:
         while process.returncode is None:
+            if max_output_lines > 0 and len(self.stderr) + len(self.stdout) > max_output_lines:
+                self._truncated = True
+                process.terminate()
             if self._terminate.is_set():
                 process.terminate()
             try:
@@ -83,7 +87,7 @@ class Cli:
     def terminate(self) -> None:
         self._terminate.set()
 
-    async def execute(self, command: str) -> Result:
+    async def execute(self, command: str, max_output_lines: int = 0) -> Result:
         self._busy = True
         c = shlex.split(command, posix=False)
         try:
@@ -92,13 +96,14 @@ class Cli:
                 self.stdout.clear()
                 self.stderr.clear()
                 self._terminate.clear()
+                self._truncated = False
                 terminated = False
                 now = datetime.now().strftime("%Y/%m/%d %H:%M:%S")
                 self.prefix_line = f"<{now}> {command}\n"
                 for terminal in self._stdout_terminals:
                     terminal.call_terminal_method("write", "\n" + self.prefix_line)
                 await asyncio.gather(
-                    self._controller(process=process),
+                    self._controller(process=process, max_output_lines=max_output_lines),
                     self._read_stdout(stream=process.stdout),
                     self._read_stderr(stream=process.stderr),
                 )
@@ -110,7 +115,9 @@ class Cli:
         finally:
             self._terminate.clear()
             self._busy = False
-        return Result(command=command, return_code=process.returncode, stdout_lines=self.stdout.copy(), stderr_lines=self.stderr.copy(), terminated=terminated)
+        return Result(
+            command=command, return_code=process.returncode, stdout_lines=self.stdout.copy(), stderr_lines=self.stderr.copy(), terminated=terminated, truncated=self._truncated
+        )
 
     async def shell(self, command: str) -> Result:
         self._busy = True

bale/interfaces/ssh.py

@@ -90,10 +90,10 @@ class Ssh(Cli):
         del self._config[self.host]
         self.write_config()
 
-    async def execute(self, command: str) -> Result:
+    async def execute(self, command: str, max_output_lines: int = 0) -> Result:
         self._base_cmd = f"{'' if self.use_key else f'sshpass -p {self.password} '} ssh -F {self._config_path} {self.host}"
         self._full_cmd = f"{self._base_cmd} {command}"
-        return await super().execute(self._full_cmd)
+        return await super().execute(self._full_cmd, max_output_lines)
 
     async def send_key(self) -> Result:
         await get_public_key(self._raw_path)

bale/interfaces/zfs.py

@@ -81,7 +81,7 @@ class Zfs:
         command = command if len(command) < 160 else command[:160] + "..."
         el.notify(command)
 
-    async def execute(self, command: str, notify: bool = True) -> Result:
+    async def execute(self, command: str, max_output_lines: int = 0, notify: bool = True) -> Result:
         if notify:
             self.notify(command)
         return Result(command=command)
@@ -244,10 +244,10 @@ class Ssh(ssh.Ssh, Zfs):
     def notify(self, command: str):
         super().notify(f"<{self.host}> {command}")
 
-    async def execute(self, command: str, notify: bool = True) -> Result:
+    async def execute(self, command: str, max_output_lines: int = 0, notify: bool = True) -> Result:
         if notify:
             self.notify(command)
-        result = await super().execute(command)
+        result = await super().execute(command, max_output_lines)
         if result.stderr != "":
             el.notify(result.stderr, type="negative")
         result.name = self.host

bale/result.py

@@ -13,6 +13,7 @@ class Result:
     stdout_lines: List[str] = field(default_factory=list)
     stderr_lines: List[str] = field(default_factory=list)
     terminated: bool = False
+    truncated: bool = False
     data: Any = None
     trace: str = ""
     cached: bool = False