Fix table search in edge cases

PiperOrigin-RevId: 432952643
PAIR-code · Mar 7, 2022 · 036283b · 036283b
1 parent 8993f9b
commit 036283b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/lit_nlp/client/elements/table.ts b/lit_nlp/client/elements/table.ts
@@ -648,8 +648,13 @@ export class DataTable extends ReactiveElement {
     // clang-format on
   }
 
+  private columnNameToId(name: string) {
+    return name.replace(/\s+/g, '');
+  }
+
   renderColumnHeader(header: ColumnHeader) {
     const title = header.name;
+    const headerId = this.columnNameToId(title);
 
     const handleBackgroundClick = (e: Event) => {
       this.resetView();
@@ -705,7 +710,8 @@ export class DataTable extends ReactiveElement {
       // Focus cursor on newly-shown input box.
       if (this.showColumnMenu) {
         window.requestAnimationFrame(() => {
-          const inputElem = this.shadowRoot!.querySelector(`th#${title} .togglable-menu-holder input`) as HTMLElement;
+          const inputElem = this.shadowRoot!.querySelector(
+              `th#${headerId} .togglable-menu-holder input`) as HTMLElement;
           inputElem.focus();
         });
       }
@@ -738,7 +744,7 @@ export class DataTable extends ReactiveElement {
 
     // clang-format off
     return html`
-        <th id=${title} @click=${handleBackgroundClick}>
+        <th id=${headerId} @click=${handleBackgroundClick}>
           <div class=${headerClasses} title=${title}>
             <div class="header-holder">
               <div @click=${toggleSort}>${header.html!}</div>

diff --git a/lit_nlp/client/modules/multilabel_module.ts b/lit_nlp/client/modules/multilabel_module.ts
@@ -235,6 +235,7 @@ export class MultilabelModule extends LitModule {
           <lit-data-table
             .columnNames=${columnNames}
             .data=${rows}
+            searchEnabled
           ></lit-data-table>
         </div>`;
   }