From 08a46e7f38acc354fbb4de27ad3be88bbb2b79a7 Mon Sep 17 00:00:00 2001
From: eight04 <eight04@gmail.com>
Date: Wed, 28 Feb 2024 16:24:10 +0800
Subject: [PATCH] Add: ByteScreen

---
 .gitignore           |  1 +
 pyte/__init__.py     |  2 +-
 pyte/screens.py      | 46 ++++++++++++++++++++++++++++++++++++++------
 tests/test_screen.py | 10 ++++++++++
 4 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/.gitignore b/.gitignore
index d0fa950..c2851c7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,3 +4,4 @@ dist/
 *.egg-info/
 .eggs/
 .pytest_cache/
+.venv/
diff --git a/pyte/__init__.py b/pyte/__init__.py
index 226e62b..571ca18 100644
--- a/pyte/__init__.py
+++ b/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
 
 
diff --git a/pyte/screens.py b/pyte/screens.py
index ebbac7c..80a67c5 100644
--- a/pyte/screens.py
+++ b/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):
+    # 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, encoding: str | None=None, **kwargs):
+        """
+        :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):
+        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)]
+
diff --git a/tests/test_screen.py b/tests/test_screen.py
index b6ba90d..140ba35 100644
--- a/tests/test_screen.py
+++ b/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"
+