Fix bug when using info_table with keys="headers"

Author: dmerejkowsky

cli_ui/__init__.py

@@ -409,17 +409,31 @@ def tabs(num: int) -> str:
 def info_table(
     data: Any, *, headers: Union[str, Sequence[str]] = (), fileobj: FileObj = sys.stdout
 ) -> None:
-    colored_data = list()
-    plain_data = list()
-    for row in data:
-        colored_row = list()
-        plain_row = list()
-        for item in row:
-            colored_str, plain_str = process_tokens(item, end="")
-            colored_row.append(colored_str)
-            plain_row.append(plain_str)
-        colored_data.append(colored_row)
-        plain_data.append(plain_row)
+    if headers == "keys":
+        colored_data: Any = dict()
+        plain_data: Any = dict()
+        for key, sequence in data.items():
+            colored_sequence = list()
+            plain_sequence = list()
+            for item in sequence:
+                colored_str, plain_str = process_tokens(item, end="")
+                colored_sequence.append(colored_str)
+                plain_sequence.append(plain_str)
+            colored_key, plain_key = process_tokens(key, end="")
+            colored_data[colored_key] = colored_sequence
+            plain_data[plain_key] = plain_sequence
+    else:
+        colored_data = list()
+        plain_data = list()
+        for row in data:
+            colored_row = list()
+            plain_row = list()
+            for item in row:
+                colored_str, plain_str = process_tokens(item, end="")
+                colored_row.append(colored_str)
+                plain_row.append(plain_str)
+            colored_data.append(colored_row)
+            plain_data.append(plain_row)
     if config_color(fileobj):
         data_for_tabulate = colored_data
     else:

cli_ui/tests/test_cli_ui.py

@@ -113,6 +113,80 @@ def test_timestamp(dumb_tty: DumbTTY, toggle_timestamp: None) -> None:
     assert datetime.datetime.strptime(match.groups()[0], "%Y-%m-%d %H:%M:%S")
 
 
+def test_table_with_lists_no_color(dumb_tty: DumbTTY) -> None:
+    headers = ["name", "score"]
+    data = [
+        [(cli_ui.bold, "John"), (cli_ui.green, 10)],
+        [(cli_ui.bold, "Jane"), (cli_ui.green, 5)],
+    ]
+    cli_ui.info_table(data, headers=headers, fileobj=dumb_tty)
+    actual = dumb_tty.getvalue()
+    # fmt: off
+    expected = (
+        "name      score\n"
+        "------  -------\n"
+        "John         10\n"
+        "Jane          5\n"
+        )
+    # fmt: on
+    assert actual == expected
+
+
+def test_table_with_dict_no_color(dumb_tty: DumbTTY) -> None:
+    data = {
+        (cli_ui.bold, "Name"): [(cli_ui.green, "Alice"), (cli_ui.green, "Bob")],
+        (cli_ui.bold, "Age"): [(cli_ui.blue, 24), (cli_ui.blue, 9)],
+    }
+    cli_ui.info_table(data, headers="keys", fileobj=dumb_tty)
+    actual = dumb_tty.getvalue()
+    # fmt: off
+    expected = (
+    "Name      Age\n"
+    "------  -----\n"
+    "Alice      24\n"
+    "Bob         9\n"
+    )
+    # fmt: on
+    assert actual == expected
+
+
+def test_table_with_dict_and_color(smart_tty: SmartTTY) -> None:
+    data = {
+        (cli_ui.bold, "Name",): [(cli_ui.green, "Alice"), (cli_ui.green, "Bob")],
+        (cli_ui.bold, "Age",): [(cli_ui.blue, 24), (cli_ui.blue, 9)],
+    }
+    cli_ui.info_table(data, headers="keys", fileobj=smart_tty)
+    actual = smart_tty.getvalue()
+    # fmt: off
+    expected = (
+    f"{BRIGHT}Name{RESET_ALL}      {BRIGHT}Age{RESET_ALL}\n"
+    "------  -----\n"
+    f"{GREEN}Alice{RESET_ALL}      {BLUE}24{RESET_ALL}\n"
+    f"{GREEN}Bob{RESET_ALL}         {BLUE}9{RESET_ALL}\n"
+    )
+    # fmt: on
+    assert actual == expected
+
+
+def test_table_with_lists_with_color(smart_tty: SmartTTY) -> None:
+    headers = ["name", "score"]
+    data = [
+        [(cli_ui.bold, "John"), (cli_ui.green, 10)],
+        [(cli_ui.bold, "Jane"), (cli_ui.green, 5)],
+    ]
+    cli_ui.info_table(data, headers=headers, fileobj=smart_tty)
+    actual = smart_tty.getvalue()
+    # fmt: off
+    expected = (
+        "name      score\n"
+        "------  -------\n"
+        f"{BRIGHT}John{RESET_ALL}         {GREEN}10{RESET_ALL}\n"
+        f"{BRIGHT}Jane{RESET_ALL}          {GREEN}5{RESET_ALL}\n"
+        )
+    # fmt: on
+    assert actual == expected
+
+
 def test_record_message(message_recorder: MessageRecorder) -> None:
     cli_ui.info_1("This is foo")
     assert message_recorder.find("foo")