From 42d189a49f9e1c7afc3a92eda4e365994ff454fc Mon Sep 17 00:00:00 2001
From: Ryan Mullins <ryanmullins@google.com>
Date: Fri, 24 Mar 2023 12:30:59 -0700
Subject: [PATCH] Adds optional column widths in Table elements via
 ColumnHeader structs

PiperOrigin-RevId: 519210840
---
 lit_nlp/client/elements/table.css                 | 15 ++++++++++++++-
 lit_nlp/client/elements/table.ts                  | 13 +++++++++++--
 lit_nlp/client/elements/table_types.ts            |  6 ++++++
 lit_nlp/client/modules/metrics_module.ts          |  8 ++++++--
 .../client/modules/salience_clustering_module.ts  |  4 ++--
 5 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/lit_nlp/client/elements/table.css b/lit_nlp/client/elements/table.css
index 947a7b8b..84541545 100644
--- a/lit_nlp/client/elements/table.css
+++ b/lit_nlp/client/elements/table.css
@@ -62,6 +62,11 @@
   padding: var(--table-cell-padding);
 }
 
+.header-holder lit-tooltip {
+  display: flex;
+  align-items: center;
+}
+
 .header-text {
   flex: 1;
   text-align: center;
@@ -73,6 +78,10 @@
   white-space: nowrap;
 }
 
+.column-header.right-align .header-text {
+  text-align: right;
+}
+
 .arrow-container {
   width: 16px;
   height: 20px;
@@ -182,7 +191,6 @@ tbody tr.focused {
 
 tbody td {
   vertical-align: top;
-  min-width: 80px;
   box-sizing: border-box;
 }
 
@@ -199,6 +207,11 @@ tbody td {
 
 .cell-holder.right-align {
   justify-content: flex-end;
+  padding-right: 24px;
+}
+
+:host([searchenabled]) .cell-holder.right-align {
+  padding-right: 46px;
 }
 
 .text-cell {
diff --git a/lit_nlp/client/elements/table.ts b/lit_nlp/client/elements/table.ts
index d1088eeb..4f5a0c43 100644
--- a/lit_nlp/client/elements/table.ts
+++ b/lit_nlp/client/elements/table.ts
@@ -658,9 +658,19 @@ export class DataTable extends ReactiveElement {
     // in the row render method
     this.selectedIndicesSetForRender = new Set<number>(this.selectedIndices);
 
+    const cols = this.columnHeaders.map((header) => {
+      const styles = styleMap({
+        'max-width': header.maxWidth || 'unset',
+        'min-width': header.minWidth || 'unset',
+        'width': header.width || 'unset',
+      });
+      return html`<col span="1" style=${styles}>`;
+    });
+
     // clang-format off
     return html`<div class="holder">
       <table class=${classMap({'has-footer': this.hasFooter})}>
+        <colgroup>${cols}</colgroup>
         <thead>
           ${this.columnHeaders.map((c, i) =>
               this.renderColumnHeader(c, i === this.columnHeaders.length - 1))}
@@ -846,12 +856,11 @@ export class DataTable extends ReactiveElement {
     const headerClasses =
         classMap({'column-header': true, 'right-align': header.rightAlign!});
 
-
     // TODO(b/255799266): Add fast tooltips to icons.
     // There's some rendering trickiness around the table element and tooltips.
     // clang-format off
     return html`
-        <th id=${headerId}>
+        <th id=${headerId} scope="col">
           <div class=${headerClasses}>
             <div class="header-holder">
               <div @click=${toggleSort}>${header.html!}</div>
diff --git a/lit_nlp/client/elements/table_types.ts b/lit_nlp/client/elements/table_types.ts
index 2edfd98f..eb3c9b73 100644
--- a/lit_nlp/client/elements/table_types.ts
+++ b/lit_nlp/client/elements/table_types.ts
@@ -51,6 +51,12 @@ export interface ColumnHeader {
    *  through selected items from the vocab list.
    */
   vocab?: string[];
+  /** The maximum width of the column, must be a valid CSS string. */
+  maxWidth?: string;
+  /** The minimum width of the column, must be a valid CSS string */
+  minWidth?: string;
+  /** The width of the column, must be a valid CSS string */
+  width?: string;
 }
 
 /** Internal data, including metadata */
diff --git a/lit_nlp/client/modules/metrics_module.ts b/lit_nlp/client/modules/metrics_module.ts
index fc306c2b..2096b2c7 100644
--- a/lit_nlp/client/modules/metrics_module.ts
+++ b/lit_nlp/client/modules/metrics_module.ts
@@ -358,8 +358,12 @@ export class MetricsModule extends LitModule {
     }
 
     const metricHeaders: ColumnHeader[] = metricNames.map(name => {
-      return {name, rightAlign: true, html: html`
-        <div class="header-text">${name}</div>`
+      const [group, metric] = name.split(': ');
+      return {
+        name,
+        rightAlign: true,
+        html: html`<div class="header-text">${group}<br>${metric}</div>`,
+        width: '100px'
       };
     });
 
diff --git a/lit_nlp/client/modules/salience_clustering_module.ts b/lit_nlp/client/modules/salience_clustering_module.ts
index 01a0dc75..883acba3 100644
--- a/lit_nlp/client/modules/salience_clustering_module.ts
+++ b/lit_nlp/client/modules/salience_clustering_module.ts
@@ -269,8 +269,8 @@ export class SalienceClusteringModule extends LitModule {
         <div class="grad-key-label">${gradKey}</div>
         <lit-data-table
           .columnNames=${[
-            'Cluster index',
-            'N',
+            {'name':'Cluster index', 'width': '125px'},
+            {'name': 'N', 'width': '75px'},
             `Tokens with high average saliency (up to ${numTokensPerCluster})`
           ]}
           .data=${rowsByClusters}