Fix drawing headerless colored tables with title (#332)

prettytable · Oct 29, 2024 · 5ffbc2d · 5ffbc2d
2 parents db520f8 + fa57608
commit 5ffbc2d
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 6 deletions.
diff --git a/src/prettytable/prettytable.py b/src/prettytable/prettytable.py
@@ -1977,8 +1977,12 @@ def get_string(self, **kwargs) -> str:
         ):
             lines.append(self._stringify_hrule(options, where="top_"))
             if title and options["vrules"] in (VRuleStyle.ALL, VRuleStyle.FRAME):
+                left_j_len = len(self.left_junction_char)
+                right_j_len = len(self.right_junction_char)
                 lines[-1] = (
-                    self.left_junction_char + lines[-1][1:-1] + self.right_junction_char
+                    self.left_junction_char
+                    + lines[-1][left_j_len:-right_j_len]
+                    + self.right_junction_char
                 )
 
         # Add rows
@@ -2000,8 +2004,10 @@ def get_string(self, **kwargs) -> str:
             lines.append(self._stringify_hrule(options, where="bottom_"))
 
         if "orgmode" in self.__dict__ and self.orgmode:
+            left_j_len = len(self.left_junction_char)
+            right_j_len = len(self.right_junction_char)
             lines = [
-                "|" + new_line[1:-1] + "|"
+                "|" + new_line[left_j_len:-right_j_len] + "|"
                 for old_line in lines
                 for new_line in old_line.split("\n")
             ]

diff --git a/tests/test_colortable.py b/tests/test_colortable.py
@@ -105,7 +105,16 @@ class TestColorTableRendering:
         Tests the color table rendering using the default alignment (`'c'`)
     """
 
-    def test_color_table_rendering(self) -> None:
+    @pytest.mark.parametrize(
+        ["with_title", "with_header"],
+        [
+            (False, True),  # the default
+            (True, True),  # titled
+            (True, False),  # titled, no header
+            (True, True),  # both title and header
+        ],
+    )
+    def test_color_table_rendering(self, with_title: bool, with_header: bool) -> None:
         """Tests the color table rendering using the default alignment (`'c'`)"""
         chars = {
             "+": "\x1b[36m+\x1b[0m\x1b[96m",
@@ -140,15 +149,23 @@ def test_color_table_rendering(self) -> None:
             (plus + minus * 3) * 6 + plus,
         )
 
-        header_str = str("\n".join(header))
-        body_str = str("\n".join(body))
+        if with_title:
+            header_str = str("\n".join(header))
+        else:
+            header_str = str(header[2])
+        if with_header:
+            body_str = str("\n".join(body))
+        else:
+            body_str = str("\n".join(body[2:]))
 
         table = ColorTable(
             ("A", "B", "C", "D", "E", "F"),
             theme=Themes.OCEAN,
         )
 
-        table.title = "Efforts"
+        if with_title:
+            table.title = "Efforts"
+        table.header = with_header
         table.add_row([1, 2, 3, 4, 5, 6])
 
         expected = header_str + "\n" + body_str + "\x1b[0m"