From dba0aaeab7be3d369d1e41b6d7cd905c04e981af Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Fri, 30 Apr 2021 18:43:38 +0545
Subject: [PATCH] accept dict for header styling (#5905)

I found that it's harder to style a header with the list using index,
they are to be applied by header name. Otherwise, we might end up using
a different header for coloring, or worse throw some error because the
header name and style don't match.
Row does have it's property of index, so it makes sense there.
---
 dvc/command/experiments.py  | 41 +++++++++++++++++--------------------
 dvc/ui/__init__.py          |  2 +-
 dvc/ui/table.py             | 19 ++++++++++-------
 tests/unit/ui/test_table.py |  1 +
 4 files changed, 33 insertions(+), 30 deletions(-)

diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py
index 344fc658c..fc0cb3cb9 100644
--- a/dvc/command/experiments.py
+++ b/dvc/command/experiments.py
@@ -362,31 +362,9 @@ def show_experiments(
         kwargs.get("sort_order"),
         kwargs.get("precision"),
     )
-    styles = [
-        {"no_wrap": True, "header_style": "black on grey93"},
-        {"header_style": "black on grey93"},
-        *[
-            {
-                "justify": "right",
-                "header_style": "black on cornsilk1",
-                "no_wrap": True,
-                "collapse": idx != 0,
-            }
-            for idx, _ in enumerate(metric_headers)
-        ],
-        *[
-            {
-                "justify": "left",
-                "header_style": "black on light_cyan1",
-                "collapse": idx != 0,
-            }
-            for idx, _ in enumerate(param_headers)
-        ],
-    ]
 
     if no_timestamp:
         td.drop("Created")
-        styles.pop(1)
 
     baseline_styler = iffy(constantly({"style": "bold"}), default={})
     row_styles = lmap(baseline_styler, td.column("is_baseline"))
@@ -396,6 +374,25 @@ def show_experiments(
     td.column("Experiment")[:] = map(prepare_exp_id, td.as_dict(merge_headers))
     td.drop(*merge_headers[1:])
 
+    headers = {"metrics": metric_headers, "params": param_headers}
+    styles = {
+        "Experiment": {"no_wrap": True, "header_style": "black on grey93"},
+        "Created": {"header_style": "black on grey93"},
+    }
+    header_bg_colors = {"metrics": "cornsilk1", "params": "light_cyan1"}
+    styles.update(
+        {
+            header: {
+                "justify": "left" if typ == "metrics" else "params",
+                "header_style": f"black on {header_bg_colors[typ]}",
+                "collapse": idx != 0,
+                "no_wrap": typ == "metrics",
+            }
+            for typ, hs in headers.items()
+            for idx, header in enumerate(hs)
+        }
+    )
+
     td.render(
         pager=pager,
         borders=True,
diff --git a/dvc/ui/__init__.py b/dvc/ui/__init__.py
index 737a0778b..a02a7e50e 100644
--- a/dvc/ui/__init__.py
+++ b/dvc/ui/__init__.py
@@ -155,7 +155,7 @@ class Console:
         rich_table: bool = False,
         force: bool = True,
         pager: bool = False,
-        header_styles: Sequence["Styles"] = None,
+        header_styles: Union[Dict[str, "Styles"], Sequence["Styles"]] = None,
         row_styles: Sequence["Styles"] = None,
         borders: Union[bool, str] = False,
     ) -> None:
diff --git a/dvc/ui/table.py b/dvc/ui/table.py
index 54dbd7509..2f2c9179f 100644
--- a/dvc/ui/table.py
+++ b/dvc/ui/table.py
@@ -1,6 +1,7 @@
+from collections import abc
 from contextlib import contextmanager
 from itertools import zip_longest
-from typing import TYPE_CHECKING, Iterator, Sequence, Union
+from typing import TYPE_CHECKING, Dict, Iterator, Sequence, Union
 
 from dvc.types import DictStrAny
 
@@ -81,7 +82,7 @@ def rich_table(
     data: TableData,
     headers: Headers = None,
     pager: bool = False,
-    header_styles: Sequence[Styles] = None,
+    header_styles: Union[Dict[str, Styles], Sequence[Styles]] = None,
     row_styles: Sequence[Styles] = None,
     borders: Union[bool, str] = False,
 ) -> None:
@@ -97,12 +98,16 @@ def rich_table(
     }
 
     table = Table(box=border_style[borders])
-    hs: Sequence[Styles] = header_styles or []
+
+    if isinstance(header_styles, abc.Sequence):
+        hs: Dict[str, Styles] = dict(zip(headers or [], header_styles))
+    else:
+        hs = header_styles or {}
+
+    for header in headers or []:
+        table.add_column(header, **hs.get(header, {}))
+
     rs: Sequence[Styles] = row_styles or []
-
-    for header, style in zip_longest(headers or [], hs):
-        table.add_column(header, **(style or {}))
-
     for row, style in zip_longest(data, rs):
         table.add_row(*row, **(style or {}))
 
diff --git a/tests/unit/ui/test_table.py b/tests/unit/ui/test_table.py
index 663175999..7087c86ba 100644
--- a/tests/unit/ui/test_table.py
+++ b/tests/unit/ui/test_table.py
@@ -119,6 +119,7 @@ def test_rich_border(capsys: CaptureFixture[str]):
     "extra_opts",
     [
         {"header_styles": [{"style": Style(bold=True)}]},
+        {"header_styles": {"first": {"style": Style(bold=True)}}},
         {"row_styles": [{"style": Style(bold=True)}]},
     ],
 )