test: add more test cases

Author: shuoweil

bigframes/display/anywidget.py

@@ -20,7 +20,7 @@
 from importlib import resources
 import functools
 import math
-from typing import Any, Dict, Iterator, List, Optional, Type
+from typing import Any, Iterator
 import uuid
 
 import pandas as pd
@@ -43,7 +43,7 @@
 except Exception:
     ANYWIDGET_INSTALLED = False
 
-WIDGET_BASE: Type[Any]
+WIDGET_BASE: type[Any]
 if ANYWIDGET_INSTALLED:
     WIDGET_BASE = anywidget.AnyWidget
 else:
@@ -65,17 +65,13 @@ class TableWidget(WIDGET_BASE):
 
     page = traitlets.Int(0).tag(sync=True)
     page_size = traitlets.Int(0).tag(sync=True)
-    row_count = traitlets.Union(
-        [traitlets.Int(), traitlets.Instance(type(None))],
-        default_value=None,
-        allow_none=True,
-    ).tag(sync=True)
+    row_count = traitlets.Int(allow_none=True, default_value=None).tag(sync=True)
     table_html = traitlets.Unicode().tag(sync=True)
     sort_column = traitlets.Unicode("").tag(sync=True)
     sort_ascending = traitlets.Bool(True).tag(sync=True)
     orderable_columns = traitlets.List(traitlets.Unicode(), []).tag(sync=True)
     _initial_load_complete = traitlets.Bool(False).tag(sync=True)
-    _batches: Optional[blocks.PandasBatches] = None
+    _batches: blocks.PandasBatches | None = None
     _error_message = traitlets.Unicode(allow_none=True, default_value=None).tag(
         sync=True
     )
@@ -88,7 +84,8 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         """
         if not ANYWIDGET_INSTALLED:
             raise ImportError(
-                "Please `pip install anywidget traitlets` or `pip install 'bigframes[anywidget]'` to use TableWidget."
+                "Please `pip install anywidget traitlets` or "
+                "`pip install 'bigframes[anywidget]'` to use TableWidget."
             )
 
         self._dataframe = dataframe
@@ -98,9 +95,9 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # Initialize attributes that might be needed by observers first
         self._table_id = str(uuid.uuid4())
         self._all_data_loaded = False
-        self._batch_iter: Optional[Iterator[pd.DataFrame]] = None
-        self._cached_batches: List[pd.DataFrame] = []
-        self._last_sort_state: Optional[_SortState] = None
+        self._batch_iter: Iterator[pd.DataFrame] | None = None
+        self._cached_batches: list[pd.DataFrame] = []
+        self._last_sort_state: _SortState | None = None
 
         # respect display options for initial page size
         initial_page_size = bigframes.options.display.max_rows
@@ -124,7 +121,10 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         self._reset_batches_for_new_page_size()
 
         if self._batches is None:
-            self._error_message = "Could not retrieve data batches. Data might be unavailable or an error occurred."
+            self._error_message = (
+                "Could not retrieve data batches. Data might be unavailable or "
+                "an error occurred."
+            )
             self.row_count = None
         elif self._batches.total_rows is None:
             # Total rows is unknown, this is an expected state.
@@ -143,7 +143,7 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         self._initial_load_complete = True
 
     @traitlets.observe("_initial_load_complete")
-    def _on_initial_load_complete(self, change: Dict[str, Any]):
+    def _on_initial_load_complete(self, change: dict[str, Any]):
         if change["new"]:
             self._set_table_html()
 
@@ -158,7 +158,7 @@ def _css(self):
         return resources.read_text(bigframes.display, "table_widget.css")
 
     @traitlets.validate("page")
-    def _validate_page(self, proposal: Dict[str, Any]) -> int:
+    def _validate_page(self, proposal: dict[str, Any]) -> int:
         """Validate and clamp the page number to a valid range.
 
         Args:
@@ -191,7 +191,7 @@ def _validate_page(self, proposal: Dict[str, Any]) -> int:
         return max(0, min(value, max_page))
 
     @traitlets.validate("page_size")
-    def _validate_page_size(self, proposal: Dict[str, Any]) -> int:
+    def _validate_page_size(self, proposal: dict[str, Any]) -> int:
         """Validate page size to ensure it's positive and reasonable.
 
         Args:
@@ -229,10 +229,6 @@ def _get_next_batch(self) -> bool:
         except StopIteration:
             self._all_data_loaded = True
             return False
-        except Exception as e:
-            # Handle other potential errors
-            self._error_message = f"Error loading data: {str(e)}"
-            return False
 
     @property
     def _batch_iterator(self) -> Iterator[pd.DataFrame]:
@@ -272,7 +268,8 @@ def _set_table_html(self) -> None:
         try:
             if self._error_message:
                 self.table_html = (
-                    f"<div class='bigframes-error-message'>{self._error_message}</div>"
+                    f"<div class='bigframes-error-message'>"
+                    f"{self._error_message}</div>"
                 )
                 return
 
@@ -297,48 +294,53 @@ def _set_table_html(self) -> None:
                 )
                 self.page = 0  # Reset to first page
 
-            start = self.page * self.page_size
-            end = start + self.page_size
-
-            # fetch more data if the requested page is outside our cache
-            cached_data = self._cached_data
-            while len(cached_data) < end and not self._all_data_loaded:
-                if self._get_next_batch():
-                    cached_data = self._cached_data
-                else:
-                    break
+            page_data = pd.DataFrame()
+            # This loop is to handle auto-correction of page number when row count is unknown
+            while True:
+                start = self.page * self.page_size
+                end = start + self.page_size
+
+                # fetch more data if the requested page is outside our cache
+                cached_data = self._cached_data
+                while len(cached_data) < end and not self._all_data_loaded:
+                    if self._get_next_batch():
+                        cached_data = self._cached_data
+                    else:
+                        break
+
+                # Get the data for the current page
+                page_data = cached_data.iloc[start:end].copy()
+
+                # Handle case where user navigated beyond available data with unknown row count
+                is_unknown_count = self.row_count is None
+                is_beyond_data = (
+                    self._all_data_loaded and len(page_data) == 0 and self.page > 0
+                )
+                if is_unknown_count and is_beyond_data:
+                    # Calculate the last valid page (zero-indexed)
+                    total_rows = len(cached_data)
+                    last_valid_page = max(0, math.ceil(total_rows / self.page_size) - 1)
+                    # Navigate back to the last valid page
+                    self.page = last_valid_page
+                    # Continue the loop to re-calculate page data
+                    continue
 
-            # Get the data for the current page
-            page_data = cached_data.iloc[start:end].copy()
+                # If page is valid, break out of the loop.
+                break
 
             # Handle index display
-            # TODO(b/438181139): Add tests for custom multiindex
             if self._dataframe._block.has_index:
-                index_name = page_data.index.name
-                page_data.insert(
-                    0, index_name if index_name is not None else "", page_data.index
+                is_unnamed_single_index = (
+                    page_data.index.name is None
+                    and not isinstance(page_data.index, pd.MultiIndex)
                 )
-            else:
-                # Default index - include as "Row" column
+                page_data = page_data.reset_index()
+                if is_unnamed_single_index and "index" in page_data.columns:
+                    page_data.rename(columns={"index": ""}, inplace=True)
+
+            # Default index - include as "Row" column if no index was present originally
+            if not self._dataframe._block.has_index:
                 page_data.insert(0, "Row", range(start + 1, start + len(page_data) + 1))
-            # Handle case where user navigated beyond available data with unknown row count
-            is_unknown_count = self.row_count is None
-            is_beyond_data = (
-                self._all_data_loaded and len(page_data) == 0 and self.page > 0
-            )
-            if is_unknown_count and is_beyond_data:
-                # Calculate the last valid page (zero-indexed)
-                total_rows = len(cached_data)
-                if total_rows > 0:
-                    last_valid_page = max(0, math.ceil(total_rows / self.page_size) - 1)
-                    # Navigate back to the last valid page
-                    self.page = last_valid_page
-                    # Recursively call to display the correct page
-                    return self._set_table_html()
-                else:
-                    # If no data at all, stay on page 0 with empty display
-                    self.page = 0
-                    return self._set_table_html()
 
             # Generate HTML table
             self.table_html = bigframes.display.html.render_html(
@@ -349,19 +351,19 @@ def _set_table_html(self) -> None:
             delattr(self, "_setting_html")
 
     @traitlets.observe("sort_column", "sort_ascending")
-    def _sort_changed(self, _change: Dict[str, Any]):
+    def _sort_changed(self, _change: dict[str, Any]):
         """Handler for when sorting parameters change from the frontend."""
         self._set_table_html()
 
     @traitlets.observe("page")
-    def _page_changed(self, _change: Dict[str, Any]) -> None:
+    def _page_changed(self, _change: dict[str, Any]) -> None:
         """Handler for when the page number is changed from the frontend."""
         if not self._initial_load_complete:
             return
         self._set_table_html()
 
     @traitlets.observe("page_size")
-    def _page_size_changed(self, _change: Dict[str, Any]) -> None:
+    def _page_size_changed(self, _change: dict[str, Any]) -> None:
         """Handler for when the page size is changed from the frontend."""
         if not self._initial_load_complete:
             return