Fixed unit tests

src/components/datagrid_new/components/DataGrid.py

@@ -4,7 +4,6 @@ import logging
 from io import BytesIO
 from typing import Literal, Any
 
-import numpy as np
 import pandas as pd
 from fasthtml.components import *
 from fasthtml.xtend import Script
@@ -127,7 +126,7 @@ class DataGrid(BaseComponent):
       columns = [DataGridColumnState(make_safe_id(col_id),
                                      col_index,
                                      col_id,
-                                     _get_column_type(_df[make_safe_id(col_id)].dtype))
+                                     _get_column_type(self._df[make_safe_id(col_id)].dtype))
                  for col_index, col_id in enumerate(_df.columns)]
       if self._state.row_index:
         columns.insert(0, DataGridColumnState(make_safe_id(ROW_INDEX_ID), -1, " ", ColumnType.RowIndex))
@@ -137,7 +136,7 @@ class DataGrid(BaseComponent):
     self._df = df.copy()
     self._df.columns = self._df.columns.map(make_safe_id)  # make sure column names are trimmed
     self._state.rows = [DataGridRowState(row_id) for row_id in self._df.index]
-    self._state.columns = _init_columns(self._df)
+    self._state.columns = _init_columns(df)  # use df not self._df to keep the original title
     self._fast_access = self._init_fast_access(self._df)
     self._total_rows = len(self._df) if self._df is not None else 0