fix mypy

Author: shuoweil

packages/bigframes/bigframes/pandas/io/api.py

@@ -300,7 +300,7 @@ def _try_read_gbq_colab_sessionless_dry_run(
 def _read_gbq_colab(  # type: ignore[overload-overlap]
     query_or_table: str,
     *,
-    callback: Optional[Callable[[bigframes.core.events.Event], None]] = ...,
+    callback: Optional[Callable[[bigframes.core.events.EventEnvelope], None]] = ...,
     pyformat_args: Optional[Dict[str, Any]] = ...,
     dry_run: Literal[False] = ...,
 ) -> bigframes.dataframe.DataFrame: ...
@@ -310,7 +310,7 @@ def _read_gbq_colab(  # type: ignore[overload-overlap]
 def _read_gbq_colab(
     query_or_table: str,
     *,
-    callback: Optional[Callable[[bigframes.core.events.Event], None]] = ...,
+    callback: Optional[Callable[[bigframes.core.events.EventEnvelope], None]] = ...,
     pyformat_args: Optional[Dict[str, Any]] = ...,
     dry_run: Literal[True] = ...,
 ) -> pandas.Series: ...
@@ -319,7 +319,7 @@ def _read_gbq_colab(
 def _read_gbq_colab(
     query_or_table: str,
     *,
-    callback: Optional[Callable[[bigframes.core.events.Event], None]] = None,
+    callback: Optional[Callable[[bigframes.core.events.EventEnvelope], None]] = None,
     pyformat_args: Optional[Dict[str, Any]] = None,
     dry_run: bool = False,
 ) -> bigframes.dataframe.DataFrame | pandas.Series:
@@ -331,7 +331,7 @@ def _read_gbq_colab(
     Args:
         query_or_table (str):
             SQL query or table ID (table ID not yet supported).
-        callback (Optional[Callable[[bigframes.core.events.Event], None]]):
+        callback (Optional[Callable[[bigframes.core.events.EventEnvelope], None]]):
             Callback to receive query execution events.
         pyformat_args (Optional[Dict[str, Any]]):
             Parameters to format into the query string.

packages/bigframes/bigframes/session/__init__.py

@@ -667,7 +667,7 @@ def _read_gbq_colab(
         self,
         query: str,
         *,
-        callback: Optional[Callable[[bigframes.core.events.Event], None]] = ...,
+        callback: Optional[Callable[[bigframes.core.events.EventEnvelope], None]] = ...,
         pyformat_args: Optional[Dict[str, Any]] = None,
         dry_run: Literal[False] = ...,
     ) -> dataframe.DataFrame: ...
@@ -677,7 +677,7 @@ def _read_gbq_colab(
         self,
         query: str,
         *,
-        callback: Optional[Callable[[bigframes.core.events.Event], None]] = ...,
+        callback: Optional[Callable[[bigframes.core.events.EventEnvelope], None]] = ...,
         pyformat_args: Optional[Dict[str, Any]] = None,
         dry_run: Literal[True] = ...,
     ) -> pandas.Series: ...
@@ -687,7 +687,9 @@ def _read_gbq_colab(
         self,
         query: str,
         *,
-        callback: Optional[Callable[[bigframes.core.events.Event], None]] = None,
+        callback: Optional[
+            Callable[[bigframes.core.events.EventEnvelope], None]
+        ] = None,
         pyformat_args: Optional[Dict[str, Any]] = None,
         dry_run: bool = False,
     ) -> Union[dataframe.DataFrame, pandas.Series]:
@@ -700,7 +702,7 @@ def _read_gbq_colab(
             query (str):
                 A SQL query string to execute. Results (if any) are turned into
                 a DataFrame.
-            callback (Optional[Callable[[bigframes.core.events.Event], None]]):
+            callback (Optional[Callable[[bigframes.core.events.EventEnvelope], None]]):
                 Callback to receive query execution events.
             pyformat_args (dict):
                 A dictionary of potential variables to replace in ``query``.