Add: ByteScreen, use_c1 option

.gitignore

@@ -4,3 +4,4 @@ dist/
 *.egg-info/
 .eggs/
 .pytest_cache/
+.venv/

pyte/__init__.py

@@ -28,7 +28,7 @@
 import io
 from typing import Union
 
-from .screens import Screen, DiffScreen, HistoryScreen, DebugScreen
+from .screens import Screen, DiffScreen, HistoryScreen, DebugScreen, ByteScreen
 from .streams import Stream, ByteStream
 
 

pyte/screens.py

@@ -47,8 +47,6 @@
 )
 from .streams import Stream
 
-wcwidth: Callable[[str], int] = lru_cache(maxsize=4096)(_wcwidth)
-
 KT = TypeVar("KT")
 VT = TypeVar("VT")
 
@@ -135,7 +133,7 @@ def __missing__(self, key: KT) -> VT:
 
 
 _DEFAULT_MODE = set([mo.DECAWM, mo.DECTCEM])
-
+_DEFAULT_WCWIDTH: Callable[[str], int] = lru_cache(maxsize=4096)(_wcwidth)
 
 class Screen:
     """
@@ -222,6 +220,7 @@ def __init__(self, columns: int, lines: int) -> None:
         self.reset()
         self.mode = _DEFAULT_MODE.copy()
         self.margins: Optional[Margins] = None
+        self.wcwidth = _DEFAULT_WCWIDTH
 
     def __repr__(self) -> str:
         return ("{0}({1}, {2})".format(self.__class__.__name__,
@@ -237,8 +236,8 @@ def render(line: StaticDefaultDict[int, Char]) -> Generator[str, None, None]:
                     is_wide_char = False
                     continue
                 char = line[x].data
-                assert sum(map(wcwidth, char[1:])) == 0
-                is_wide_char = wcwidth(char[0]) == 2
+                assert sum(map(self.wcwidth, char[1:])) == 0
+                is_wide_char = self.wcwidth(char[0]) == 2
                 yield char
 
         return ["".join(render(self.buffer[y])) for y in range(self.lines)]
@@ -479,7 +478,7 @@ def draw(self, data: str) -> None:
             self.g1_charset if self.charset else self.g0_charset)
 
         for char in data:
-            char_width = wcwidth(char)
+            char_width = self.wcwidth(char)
 
             # If this was the last column in a line and auto wrap mode is
             # enabled, move the cursor to the beginning of the next line,
@@ -1337,3 +1336,38 @@ def __getattribute__(self, attr: str) -> Callable[..., None]:
             return self.only_wrapper(attr)
         else:
             return lambda *args, **kwargs: None
+
+def byte_screen_wcwidth(text: str) -> int:
+    # FIXME: should we always return 1?
+    n = _DEFAULT_WCWIDTH(text)
+    if n <= 0 and text <= "\xff":
+        return 1
+    return n
+
+class ByteScreen(Screen):
+    """A screen that draws bytes and stores byte-string in the buffer, including un-printable/zero-length chars."""
+    def __init__(self, *args: Any, encoding: str|None = None, **kwargs: Any):
+        """
+        :param encoding: The encoding of the screen. If set, the byte-string will be decoded when calling :meth:`ByteScreen.display`.
+        """
+        super().__init__(*args, **kwargs)
+        self.encoding = encoding
+        self.wcwidth = byte_screen_wcwidth
+
+    def draw(self, data: str|bytes) -> None:
+        if isinstance(data, bytes):
+            data = data.decode("latin-1")
+        return super().draw(data)
+
+    @property
+    def display(self) -> List[str]:
+        if not self.encoding:
+            return super().display
+
+        def render(line: StaticDefaultDict[int, Char]) -> Generator[str, None, None]:
+            for x in range(self.columns):
+                char = line[x].data
+                yield char
+
+        return ["".join(render(self.buffer[y])).encode("latin-1").decode(self.encoding) for y in range(self.lines)]
+

pyte/streams.py

@@ -139,10 +139,11 @@ class Stream:
         "[^" + "".join(map(re.escape, _special)) + "]+")
     del _special
 
-    def __init__(self, screen: Optional[Screen] = None, strict: bool = True) -> None:
+    def __init__(self, screen: Optional[Screen] = None, strict: bool = True, use_c1: bool = True) -> None:
         self.listener: Optional[Screen] = None
         self.strict = strict
         self.use_utf8: bool = True
+        self.use_c1: bool = use_c1
 
         self._taking_plain_text: Optional[bool] = None
 
@@ -237,6 +238,9 @@ def _parser_fsm(self) -> ParserGenerator:
         draw = listener.draw
         debug = listener.debug
 
+        use_c1 = self.use_c1
+        in_control = False
+
         ESC, CSI_C1 = ctrl.ESC, ctrl.CSI_C1
         OSC_C1 = ctrl.OSC_C1
         SP_OR_GT = ctrl.SP + ">"
@@ -274,6 +278,7 @@ def create_dispatcher(mapping: Mapping[str, str]) -> Dict[str, Callable[..., Non
                 #    recognizes ``ESC % C`` sequences for selecting control
                 #    character set. However, in the current version these
                 #    are noop.
+                in_control = True
                 char = yield None
                 if char == "[":
                     char = CSI_C1  # Go to CSI.
@@ -304,7 +309,7 @@ def create_dispatcher(mapping: Mapping[str, str]) -> Dict[str, Callable[..., Non
                     continue
 
                 basic_dispatch[char]()
-            elif char == CSI_C1:
+            elif char == CSI_C1 and (use_c1 or in_control):
                 # All parameters are unsigned, positive decimal integers, with
                 # the most significant digit sent first. Any parameter greater
                 # than 9999 is set to 9999. If you do not specify a value, a 0
@@ -354,11 +359,14 @@ def create_dispatcher(mapping: Mapping[str, str]) -> Dict[str, Callable[..., Non
                             else:
                                 csi_dispatch[char](*params)
                             break  # CSI is finished.
-            elif char == OSC_C1:
+                in_control = False
+            elif char == OSC_C1 and (use_c1 or in_control):
                 code = yield None
                 if code == "R":
+                    in_control = False
                     continue  # Reset palette. Not implemented.
                 elif code == "P":
+                    in_control = False
                     continue  # Set palette. Not implemented.
 
                 param = ""
@@ -376,6 +384,8 @@ def create_dispatcher(mapping: Mapping[str, str]) -> Dict[str, Callable[..., Non
                     listener.set_icon_name(param)
                 if code in "02":
                     listener.set_title(param)
+
+                in_control = False
             elif char not in NUL_OR_DEL:
                 draw(char)
 

tests/test_screen.py

@@ -1583,3 +1583,13 @@ def test_screen_set_icon_name_title():
 
     screen.set_title(text)
     assert screen.title == text
+
+
+def test_byte_screen() -> None:
+    screen = pyte.ByteScreen(10, 1, encoding="big5")
+
+    text = "限".encode("big5")
+    screen.draw(text)
+    assert screen.display[0].strip() == "限"
+    assert screen.buffer[0][0].data == "\xad"
+

tests/test_stream.py

@@ -332,3 +332,22 @@ def test_byte_stream_select_other_charset():
     # c) enable utf-8
     stream.select_other_charset("G")
     assert stream.use_utf8
+
+
+def test_byte_stream_without_c1() -> None:
+    screen = pyte.ByteScreen(3, 3)
+    stream = pyte.ByteStream(screen, use_c1=False)
+    stream.select_other_charset("@")
+    b = '𡶐'.encode("big5-hkscs")
+    stream.feed(b)
+    assert screen.display[0] == "\x9b\xf3 "
+    assert stream.use_c1 == False
+
+def test_byte_stream_without_c1_with_c0() -> None:
+    screen = pyte.ByteScreen(3, 3)
+    stream = pyte.ByteStream(screen, use_c1=False)
+    stream.select_other_charset("@")
+    stream.feed(b"\x1b[1;36m")
+    assert screen.display[0] == "   "
+
+