Fix decoding error (issue 168)

aexpect/client.py

@@ -27,6 +27,7 @@
 import subprocess
 import threading
 import time
+from codecs import getincrementaldecoder
 
 from aexpect.exceptions import (
     ExpectError,
@@ -56,6 +57,9 @@
 
 LOG = logging.getLogger(__name__)
 
+# Buffer size in byte for pipe reads
+READ_BUFFER_SIZE = 1024
+
 
 def kill_tail_threads():
     """
@@ -731,6 +735,8 @@ def _print_line(text):
             poller = select.poll()
             poller.register(tail_pipe, select.POLLIN)
             bfr = ""
+            decoder_class = getincrementaldecoder(self.encoding)
+            decoder = decoder_class(errors="ignore")
             while True:
                 if _THREAD_KILL_REQUESTED.is_set():
                     try:
@@ -745,10 +751,10 @@ def _print_line(text):
                     break
                 if poll_status:
                     # Some data is available; read it
-                    new_data = os.read(tail_pipe, 1024)
-                    if not new_data:
+                    new_bytes = os.read(tail_pipe, READ_BUFFER_SIZE)
+                    if not new_bytes:
                         break
-                    new_data = new_data.decode(self.encoding, "ignore")
+                    new_data = decoder.decode(input=new_bytes)
                     if not new_data:  # all chars were ignored, skip round
                         continue
                     bfr += new_data
@@ -901,23 +907,23 @@ def _read_nonblocking(self, internal_timeout=None, timeout=None):
         expect_pipe = self._get_fd("expect")
         poller = select.poll()
         poller.register(expect_pipe, select.POLLIN)
-        data = ""
+        data = b""
         read = 0
         while True:
             try:
                 poll_status = poller.poll(internal_timeout)
             except select.error:
                 return read, data
             if poll_status:
-                raw_data = os.read(expect_pipe, 1024)
+                raw_data = os.read(expect_pipe, READ_BUFFER_SIZE)
                 if not raw_data:
-                    return read, data
+                    return read, data.decode(self.encoding, "ignore")
                 read += len(raw_data)
-                data += raw_data.decode(self.encoding, "ignore")
+                data += raw_data
             else:
-                return read, data
+                return read, data.decode(self.encoding, "ignore")
             if end_time and time.monotonic() > end_time:
-                return read, data
+                return read, data.decode(self.encoding, "ignore")
 
     def read_nonblocking(self, internal_timeout=None, timeout=None):
         """

tests/test_client.py

@@ -18,6 +18,7 @@
 import random
 import string
 import sys
+import time
 import unittest
 
 from aexpect import client
@@ -194,5 +195,156 @@ def get_proc_fds():
         )
 
 
+class EncodingTest(unittest.TestCase):
+
+    DEBUG = False
+
+    # Encoding used to translate between Unicode text and bytes
+    ENCODING = "utf-8"
+
+    # text whose characters decode to multiple byte
+    TEXT = "嗨😀"
+
+    REPETITIONS_FOR_TAIL = 3
+
+    MAX_OFFSET = 10
+
+    def analyze_output(self, offset, new_output):
+        """Helper; Compare output to expectation"""
+        # remove the leading offset whitespace
+        idx = 0
+        for idx, char in enumerate(new_output):
+            if char.isspace():
+                continue
+            if char == self.TEXT[0]:
+                break
+            self.fail(
+                f"Unexpected char found at {idx=}: {char!r} ({char.encode(self.ENCODING)}). "
+                f"Line start: {new_output[:50]}, line length: {len(new_output)}"
+            )
+        if idx == len(new_output):
+            self.fail("Test text not found!")
+        if idx > 0:
+            new_output = new_output[idx:]
+        if self.DEBUG:
+            print(f"Skipping {idx} whitespace chars at start")
+
+        # print start and end, count chars
+        n_chars = len(new_output)
+        if self.DEBUG:
+            print(f"Output for offset {offset}: len={n_chars}.")
+            for idx, char in enumerate(new_output[:3]):
+                print_char = chr(0x21B2) if char == "\n" else char
+                print(
+                    f"char {idx}: {print_char} ({char.encode(self.ENCODING)})",
+                    end="; ",
+                )
+            print("...", end="")
+            for idx, char in enumerate(new_output[-3:]):
+                print_char = chr(0x21B2) if char == "\n" else char
+                print(
+                    f"char {n_chars-3+idx}: {print_char} ({char.encode(self.ENCODING)})",
+                    end="; ",
+                )
+            print()
+        return n_chars
+
+    def analyze_results(self, all_lengths: list):
+        """Helper: compare results, decide whether test was successful"""
+        if not all_lengths:
+            self.fail("no successful output analyses")
+        expect = all_lengths[0]
+        if any(curr_length != expect for curr_length in all_lengths[1:]):
+            self.fail("There were differences in encoded output lengths")
+        elif self.DEBUG:
+            print("SUCCESS")
+
+    @unittest.skipUnless(os.name == "posix", "Unix/Linux/macOS only")
+    def test_shell(self):
+        """
+        Tests correct decoding of multibyte characters in ShellSession.
+
+        Even if reading is interrupted with incomplete characters, we
+        expect correct output.
+
+        Spawns a python session that produces multibyte output
+        with various single-byte offsets.
+        """
+        sess = client.ShellSession("/bin/sh")
+        sess.cmd_output(
+            "echo 'Just removing potential initial prompt from output'"
+        )
+        all_lengths = []
+        output = self.TEXT.encode(self.ENCODING)
+        repetitions = 1024 // len(output) + 1
+        for offset in range(self.MAX_OFFSET):
+            if self.DEBUG:
+                print(f"Start testing with shell and offset {offset}")
+            cmd = (
+                f"import os; import sys; t=b' '*{offset}+{output!r}*{repetitions}+b'\\n'; "
+                f"f=os.fdopen(sys.stdout.fileno(), 'wb', closefd=False); f.write(t); f.flush()"
+            )
+            new_output = sess.cmd_output(f'{sys.executable} -c "{cmd}"')
+            all_lengths.append(self.analyze_output(offset, new_output))
+        sess.close()
+        self.analyze_results(all_lengths)
+
+    def test_tail(self):
+        """
+        Tests correct decoding of multibyte characters in Tail.
+
+        Like test_shell, but using a Tail and repeating the output to get
+        multiple lines of output. Requires custom output gatherer and
+        termination function
+        """
+        output_buffer = []
+        terminated = False
+
+        def remember_output(new_output):
+            nonlocal output_buffer
+            output_buffer.append(new_output)
+
+        def termination_func(_status):
+            nonlocal terminated
+            terminated = True
+
+        output = self.TEXT.encode(self.ENCODING)
+        repetitions = 1024 // len(output) + 1
+        all_lengths = []
+        for offset in range(self.MAX_OFFSET):
+            terminated = False
+            output_buffer = []
+            cmd = (
+                f"import os; import sys; t=b' '*{offset}+{output!r}*{repetitions}+b'\\n';"
+                f"f=os.fdopen(sys.stdout.fileno(), 'wb', closefd=False); f.write(t); f.flush()"
+            )
+            for _ in range(self.REPETITIONS_FOR_TAIL - 1):
+                cmd += "; f.write(t); f.flush()"
+            if self.DEBUG:
+                print("Spawning Tail")
+            python = client.Tail(
+                f'{sys.executable} -c "{cmd}"',
+                output_func=remember_output,
+                termination_func=termination_func,
+            )
+            if self.DEBUG:
+                print(f"Listening for subproc {python.get_pid()}")
+            for _ in range(1000):
+                if terminated:
+                    break
+                if self.DEBUG:
+                    print(".", end="", flush=True)
+                time.sleep(0.01)
+            if self.DEBUG:
+                print("\nDone")
+            python.close()
+            for line in output_buffer:
+                if line.startswith("(Process terminated "):
+                    continue
+                all_lengths.append(self.analyze_output(offset, line))
+
+        self.analyze_results(all_lengths)
+
+
 if __name__ == "__main__":
     unittest.main()