fix: Respect display.progress_bar=None in background threads

packages/bigframes/bigframes/core/events.py

@@ -18,17 +18,21 @@
 import datetime
 import threading
 import uuid
-from typing import Any, Callable, Optional, Set
+from typing import Any, Callable, Literal, Set
 
 import google.cloud.bigquery._job_helpers
 import google.cloud.bigquery.job.query
 import google.cloud.bigquery.table
 
 import bigframes.session.executor
 
+_DEFAULT: Literal["default"] = "default"
+
 
 class Subscriber:
-    def __init__(self, callback: Callable[[Event], None], *, publisher: Publisher):
+    def __init__(
+        self, callback: Callable[[Event], None], *, publisher: Publisher
+    ):  # noqa: E501
         self._publisher = publisher
         self._callback = callback
         self._subscriber_id = uuid.uuid4()
@@ -81,16 +85,22 @@ def unsubscribe(self, subscriber: Subscriber):
         with self._subscribers_lock:
             self._subscribers.remove(subscriber)
 
-    def publish(self, event: Event):
+    def publish(self, envelope: "EventEnvelope"):
         with self._subscribers_lock:
             for subscriber in self._subscribers:
-                subscriber(event)
+                subscriber(envelope)
 
 
 class Event:
     pass
 
 
+@dataclasses.dataclass(frozen=True)
+class EventEnvelope:
+    event: Event
+    progress_bar: Literal["default", "auto", "notebook", "terminal"] | None = None
+
+
 @dataclasses.dataclass(frozen=True)
 class SessionClosed(Event):
     session_id: str
@@ -106,7 +116,7 @@ class ExecutionRunning(Event):
 
 @dataclasses.dataclass(frozen=True)
 class ExecutionFinished(Event):
-    result: Optional[bigframes.session.executor.ExecuteResult] = None
+    result: bigframes.session.executor.ExecuteResult | None = None
 
 
 @dataclasses.dataclass(frozen=True)
@@ -121,13 +131,16 @@ class BigQuerySentEvent(ExecutionRunning):
     """Query sent to BigQuery."""
 
     query: str
-    billing_project: Optional[str] = None
-    location: Optional[str] = None
-    job_id: Optional[str] = None
-    request_id: Optional[str] = None
+    billing_project: str | None = None
+    location: str | None = None
+    job_id: str | None = None
+    request_id: str | None = None
 
     @classmethod
-    def from_bqclient(cls, event: google.cloud.bigquery._job_helpers.QuerySentEvent):
+    def from_bqclient(
+        cls,
+        event: google.cloud.bigquery._job_helpers.QuerySentEvent,
+    ):
         return cls(
             query=event.query,
             billing_project=event.billing_project,
@@ -142,13 +155,16 @@ class BigQueryRetryEvent(ExecutionRunning):
     """Query sent another time because the previous attempt failed."""
 
     query: str
-    billing_project: Optional[str] = None
-    location: Optional[str] = None
-    job_id: Optional[str] = None
-    request_id: Optional[str] = None
+    billing_project: str | None = None
+    location: str | None = None
+    job_id: str | None = None
+    request_id: str | None = None
 
     @classmethod
-    def from_bqclient(cls, event: google.cloud.bigquery._job_helpers.QueryRetryEvent):
+    def from_bqclient(
+        cls,
+        event: google.cloud.bigquery._job_helpers.QueryRetryEvent,
+    ):
         return cls(
             query=event.query,
             billing_project=event.billing_project,
@@ -162,19 +178,20 @@ def from_bqclient(cls, event: google.cloud.bigquery._job_helpers.QueryRetryEvent
 class BigQueryReceivedEvent(ExecutionRunning):
     """Query received and acknowledged by the BigQuery API."""
 
-    billing_project: Optional[str] = None
-    location: Optional[str] = None
-    job_id: Optional[str] = None
-    statement_type: Optional[str] = None
-    state: Optional[str] = None
-    query_plan: Optional[list[google.cloud.bigquery.job.query.QueryPlanEntry]] = None
-    created: Optional[datetime.datetime] = None
-    started: Optional[datetime.datetime] = None
-    ended: Optional[datetime.datetime] = None
+    billing_project: str | None = None
+    location: str | None = None
+    job_id: str | None = None
+    statement_type: str | None = None
+    state: str | None = None
+    query_plan: list[google.cloud.bigquery.job.query.QueryPlanEntry] | None = None
+    created: datetime.datetime | None = None
+    started: datetime.datetime | None = None
+    ended: datetime.datetime | None = None
 
     @classmethod
     def from_bqclient(
-        cls, event: google.cloud.bigquery._job_helpers.QueryReceivedEvent
+        cls,
+        event: google.cloud.bigquery._job_helpers.QueryReceivedEvent,
     ):
         return cls(
             billing_project=event.billing_project,
@@ -193,21 +210,22 @@ def from_bqclient(
 class BigQueryFinishedEvent(ExecutionRunning):
     """Query finished successfully."""
 
-    billing_project: Optional[str] = None
-    location: Optional[str] = None
-    query_id: Optional[str] = None
-    job_id: Optional[str] = None
-    destination: Optional[google.cloud.bigquery.table.TableReference] = None
-    total_rows: Optional[int] = None
-    total_bytes_processed: Optional[int] = None
-    slot_millis: Optional[int] = None
-    created: Optional[datetime.datetime] = None
-    started: Optional[datetime.datetime] = None
-    ended: Optional[datetime.datetime] = None
+    billing_project: str | None = None
+    location: str | None = None
+    query_id: str | None = None
+    job_id: str | None = None
+    destination: google.cloud.bigquery.table.TableReference | None = None
+    total_rows: int | None = None
+    total_bytes_processed: int | None = None
+    slot_millis: int | None = None
+    created: datetime.datetime | None = None
+    started: datetime.datetime | None = None
+    ended: datetime.datetime | None = None
 
     @classmethod
     def from_bqclient(
-        cls, event: google.cloud.bigquery._job_helpers.QueryFinishedEvent
+        cls,
+        event: google.cloud.bigquery._job_helpers.QueryFinishedEvent,
     ):
         return cls(
             billing_project=event.billing_project,

packages/bigframes/bigframes/formatting_helpers.py

@@ -71,7 +71,10 @@ def repr_query_job(query_job: Optional[bigquery.QueryJob]):
     if query_job is None:
         return "No job information available"
     if query_job.dry_run:
-        return f"Computation deferred. Computation will process {get_formatted_bytes(query_job.total_bytes_processed)}"
+        return (
+            f"Computation deferred. Computation will process "
+            f"{get_formatted_bytes(query_job.total_bytes_processed)}"
+        )
     res = "Query Job Info"
     for key, value in query_job_prop_pairs.items():
         job_val = getattr(query_job, value)
@@ -105,11 +108,15 @@ def repr_query_job_html(query_job: Optional[bigquery.QueryJob]):
     if query_job is None:
         return "No job information available"
     if query_job.dry_run:
-        return f"Computation deferred. Computation will process {get_formatted_bytes(query_job.total_bytes_processed)}"
+        return (
+            f"Computation deferred. Computation will process "
+            f"{get_formatted_bytes(query_job.total_bytes_processed)}"
+        )
 
     # We can reuse the plaintext repr for now or make a nicer table.
-    # For deferred mode consistency, let's just wrap the text in a pre block or similar,
-    # but the request implies we want a distinct HTML representation if possible.
+    # For deferred mode consistency, let's just wrap the text in a pre
+    # block or similar, but the request implies we want a distinct HTML
+    # representation if possible.
     # However, existing repr_query_job returns a simple string.
     # Let's format it as a simple table or list.
 
@@ -123,7 +130,10 @@ def repr_query_job_html(query_job: Optional[bigquery.QueryJob]):
                     location=query_job.location,
                     job_id=query_job.job_id,
                 )
-                res += f'<li>Job: <a target="_blank" href="{url}">{query_job.job_id}</a></li>'
+                res += (
+                    f'<li>Job: <a target="_blank" href="{url}">'
+                    f"{query_job.job_id}</a></li>"
+                )
             elif key == "Slot Time":
                 res += f"<li>{key}: {get_formatted_time(job_val)}</li>"
             elif key == "Bytes Processed":
@@ -138,7 +148,7 @@ def repr_query_job_html(query_job: Optional[bigquery.QueryJob]):
 
 
 def progress_callback(
-    event: bigframes.core.events.Event,
+    envelope: Any,
 ):
     """Displays a progress bar while the query is running"""
     global current_display_id
@@ -147,12 +157,21 @@ def progress_callback(
         import bigframes._config
         import bigframes.core.events
     except ImportError:
-        # Since this gets called from __del__, skip if the import fails to avoid
+        # Since this gets called from __del__, skip if the import fails
+        # to avoid
         # ImportError: sys.meta_path is None, Python is likely shutting down.
         # This will allow cleanup to continue.
         return
 
-    progress_bar = bigframes._config.options.display.progress_bar
+    if isinstance(envelope, bigframes.core.events.EventEnvelope):
+        event = envelope.event
+        progress_bar = envelope.progress_bar
+    else:
+        event = envelope
+        progress_bar = bigframes.core.events._DEFAULT
+
+    if progress_bar == bigframes.core.events._DEFAULT:
+        progress_bar = bigframes._config.options.display.progress_bar
 
     if progress_bar == "auto":
         progress_bar = "notebook" if in_ipython() else "terminal"
@@ -232,7 +251,8 @@ def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None):
             job.result()
             job.reload()
             display.update_display(
-                display.HTML(get_base_job_loading_html(job)), display_id=display_id
+                display.HTML(get_base_job_loading_html(job)),
+                display_id=display_id,
             )
         elif progress_bar == "terminal":
             inital_loading_bar = get_base_job_loading_string(job)
@@ -286,7 +306,10 @@ def render_job_link_html(
         job_id=job_id,
     )
     if job_url:
-        job_link = f' [<a target="_blank" href="{job_url}">Job {project_id}:{location}.{job_id} details</a>]'
+        job_link = (
+            f' [<a target="_blank" href="{job_url}">'
+            f"Job {project_id}:{location}.{job_id} details</a>]"
+        )
     else:
         job_link = ""
     return job_link
@@ -323,7 +346,10 @@ def get_job_url(
     """
     if project_id is None or location is None or job_id is None:
         return None
-    return f"""https://console.cloud.google.com/bigquery?project={project_id}&j=bq:{location}:{job_id}&page=queryresults"""
+    return (
+        f"https://console.cloud.google.com/bigquery?project={project_id}"
+        f"&j=bq:{location}:{job_id}&page=queryresults"
+    )
 
 
 def render_bqquery_sent_event_html(
@@ -348,7 +374,10 @@ def render_bqquery_sent_event_html(
         job_id=event.job_id,
         request_id=event.request_id,
     )
-    query_text_details = f"<details><summary>SQL</summary><pre>{html.escape(event.query)}</pre></details>"
+    query_text_details = (
+        f"<details><summary>SQL</summary><pre>"
+        f"{html.escape(event.query)}</pre></details>"
+    )
 
     return f"""
     Query started{query_id}.{job_link}{query_text_details}
@@ -397,7 +426,10 @@ def render_bqquery_retry_event_html(
         job_id=event.job_id,
         request_id=event.request_id,
     )
-    query_text_details = f"<details><summary>SQL</summary><pre>{html.escape(event.query)}</pre></details>"
+    query_text_details = (
+        f"<details><summary>SQL</summary><pre>"
+        f"{html.escape(event.query)}</pre></details>"
+    )
 
     return f"""
     Retrying query{query_id}.{job_link}{query_text_details}
@@ -443,7 +475,10 @@ def render_bqquery_received_event_html(
     query_plan_details = ""
     if event.query_plan:
         plan_str = "\n".join([str(entry) for entry in event.query_plan])
-        query_plan_details = f"<details><summary>Query Plan</summary><pre>{html.escape(plan_str)}</pre></details>"
+        query_plan_details = (
+            f"<details><summary>Query Plan</summary><pre>"
+            f"{html.escape(plan_str)}</pre></details>"
+        )
 
     return f"""
     Query{query_id} is {event.state}.{job_link}{query_plan_details}
@@ -506,7 +541,8 @@ def render_bqquery_finished_event_plaintext(
 
     bytes_str = ""
     if event.total_bytes_processed is not None:
-        bytes_str = f" {humanize.naturalsize(event.total_bytes_processed)} processed."
+        size_str = humanize.naturalsize(event.total_bytes_processed)
+        bytes_str = f" {size_str} processed."
 
     slot_time_str = ""
     if event.slot_millis is not None:
@@ -572,7 +608,8 @@ def get_formatted_time(val):
         Duration string
     """
     try:
-        return humanize.naturaldelta(datetime.timedelta(milliseconds=float(val)))
+        delta = datetime.timedelta(milliseconds=float(val))
+        return humanize.naturaldelta(delta)
     except Exception:
         return val
 
@@ -591,7 +628,10 @@ def get_formatted_bytes(val):
 
 
 def get_bytes_processed_string(val: Any):
-    """Try to get bytes processed string. Return empty if passed non int value"""
+    """Try to get bytes processed string.
+
+    Return empty if passed non int value.
+    """
     bytes_processed_string = ""
     if isinstance(val, int):
         bytes_processed_string = f"""{get_formatted_bytes(val)} processed. """

packages/bigframes/bigframes/functions/function.py

@@ -162,7 +162,8 @@ class Udf(Protocol):
     """
 
     @property
-    def udf_def(self) -> udf_def.BigqueryUdf: ...
+    def udf_def(self) -> udf_def.BigqueryUdf:
+        ...
 
 
 class BigqueryCallableRoutine: