Merge pull request #58 from TankerHQ/dm/info-table-headers

Fix using info_table with `keys="headers"`

Author: dmerejkowsky

cli_ui/__init__.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union, IO
+from typing import IO, Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
 
 Dict, Type
 
@@ -15,8 +15,8 @@
 import traceback
 
 import colorama
-import unidecode
 import tabulate
+import unidecode
 
 ConfigValue = Union[None, bool, str]
 FileObj = IO[str]
@@ -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/conftest.py

@@ -1,9 +1,10 @@
-from typing import Any, Iterator, Optional
-import cli_ui
 import re
+from typing import Any, Iterator, Optional
 
 import pytest
 
+import cli_ui
+
 
 class MessageRecorder:
     """ Helper class to tests emitted messages """

cli_ui/tests/test_cli_ui.py

@@ -1,20 +1,22 @@
-from typing import Iterator
 import datetime
 import io
 import operator
 import re
+from typing import Iterator
 from unittest import mock
 
-import colorama.ansi
 import colorama
+import colorama.ansi
 import pytest
 
 import cli_ui
 from cli_ui.tests.conftest import MessageRecorder
 
-
-def assert_equal_strings(a: str, b: str) -> bool:
-    return a.split() == b.split()
+BLUE = colorama.Fore.BLUE
+GREEN = colorama.Fore.GREEN
+RED = colorama.Fore.RED
+RESET_ALL = colorama.Style.RESET_ALL
+BRIGHT = colorama.Style.BRIGHT
 
 
 class SmartTTY(io.StringIO):
@@ -59,16 +61,10 @@ def test_info_stdout_is_a_tty(smart_tty: io.StringIO) -> None:
     )
     # fmt: on
     expected = (
-        colorama.Fore.RED
-        + "this is red "
-        + colorama.Style.RESET_ALL
-        + colorama.Fore.GREEN
-        + "this is green"
-        + colorama.Style.RESET_ALL
-        + "\n"
+        RED + "this is red " + RESET_ALL + GREEN + "this is green" + "\n" + RESET_ALL
     )
     actual = smart_tty.getvalue()
-    assert_equal_strings(actual, expected)
+    assert actual == expected
 
 
 def test_update_title(smart_tty: SmartTTY) -> None:
@@ -78,17 +74,13 @@ def test_update_title(smart_tty: SmartTTY) -> None:
         fileobj=smart_tty,
         update_title=True
     )
-    # fmt: on
     expected = (
-        colorama.ansi.set_title("Something bold")
-        + "Something "
-        + colorama.Style.BRIGHT
-        + "bold"
-        + colorama.Style.RESET_ALL
-        + "\n"
+        "\x1b]0;Something bold\n\x07"
+        f"Something {BRIGHT}bold\n{RESET_ALL}"
     )
+    # fmt: on
     actual = smart_tty.getvalue()
-    assert_equal_strings(actual, expected)
+    assert actual == expected
 
 
 def test_info_stdout_is_not_a_tty(dumb_tty: DumbTTY) -> None:
@@ -101,23 +93,16 @@ def test_info_stdout_is_not_a_tty(dumb_tty: DumbTTY) -> None:
     # fmt: on
     expected = "this is red this is green\n"
     actual = dumb_tty.getvalue()
-    assert_equal_strings(actual, expected)
+    assert actual == expected
 
 
 def test_info_characters(smart_tty: SmartTTY) -> None:
     cli_ui.info(
         "Doing stuff", cli_ui.ellipsis, "sucess", cli_ui.check, fileobj=smart_tty
     )
     actual = smart_tty.getvalue()
-    expected = (
-        "Doing stuff "
-        + colorama.Style.RESET_ALL
-        + "…"
-        + " sucess "
-        + colorama.Fore.GREEN
-        + "✓"
-    )
-    assert_equal_strings(actual, expected)
+    expected = f"Doing stuff {RESET_ALL}{RESET_ALL}… {RESET_ALL}sucess {RESET_ALL}{GREEN}✓ {RESET_ALL}\n{RESET_ALL}"
+    assert actual == expected
 
 
 def test_timestamp(dumb_tty: DumbTTY, toggle_timestamp: None) -> None:
@@ -128,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")

docs/changelog.rst

@@ -1,6 +1,11 @@
 Changelog
 ----------
 
+v0.12.0
++++++++
+
+* Fix using `info_table` with `keys="headers"`
+
 v0.11.0
 +++++++
 

lint.sh

@@ -2,6 +2,7 @@ set -x
 set -e
 
 poetry run black --check .
+poetry run isort --check .
 poetry run mypy cli_ui
 poetry run pyflakes cli_ui/__init__.py cli_ui/tests/test_cli_ui.py
 poetry run sphinx-build -W docs docs/_build/html