Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit04d0273

Browse files
feat: add progress bar for magics (#396)
* feat: add progress bar for magics* feat: remove default progress bar* feat: add default tqdm value in magic
1 parent53dff2a commit04d0273

File tree

3 files changed

+107
-7
lines changed

3 files changed

+107
-7
lines changed

‎google/cloud/bigquery/_tqdm_helpers.py‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,15 +55,14 @@ def get_progress_bar(progress_bar_type, description, total, unit):
5555

5656
defwait_for_query(query_job,progress_bar_type=None):
5757
"""Return query result and display a progress bar while the query running, if tqdm is installed."""
58-
ifprogress_bar_typeisNone:
59-
returnquery_job.result()
60-
6158
default_total=1
6259
current_stage=None
6360
start_time=time.time()
6461
progress_bar=get_progress_bar(
6562
progress_bar_type,"Query is running",default_total,"query"
6663
)
64+
ifprogress_barisNone:
65+
returnquery_job.result()
6766
i=0
6867
whileTrue:
6968
ifquery_job.query_plan:

‎google/cloud/bigquery/magics/magics.py‎

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,7 @@ def __init__(self):
182182
self._default_query_job_config=bigquery.QueryJobConfig()
183183
self._bigquery_client_options=client_options.ClientOptions()
184184
self._bqstorage_client_options=client_options.ClientOptions()
185+
self._progress_bar_type="tqdm"
185186

186187
@property
187188
defcredentials(self):
@@ -313,6 +314,26 @@ def default_query_job_config(self):
313314
defdefault_query_job_config(self,value):
314315
self._default_query_job_config=value
315316

317+
@property
318+
defprogress_bar_type(self):
319+
"""str: Default progress bar type to use to display progress bar while
320+
executing queries through IPython magics.
321+
322+
Note::
323+
Install the ``tqdm`` package to use this feature.
324+
325+
Example:
326+
Manually setting the progress_bar_type:
327+
328+
>>> from google.cloud.bigquery import magics
329+
>>> magics.context.progress_bar_type = "tqdm"
330+
"""
331+
returnself._progress_bar_type
332+
333+
@progress_bar_type.setter
334+
defprogress_bar_type(self,value):
335+
self._progress_bar_type=value
336+
316337

317338
context=Context()
318339

@@ -524,6 +545,15 @@ def _create_dataset_if_necessary(client, dataset_id):
524545
"name (ex. $my_dict_var)."
525546
),
526547
)
548+
@magic_arguments.argument(
549+
"--progress_bar_type",
550+
type=str,
551+
default=None,
552+
help=(
553+
"Sets progress bar type to display a progress bar while executing the query."
554+
"Defaults to use tqdm. Install the ``tqdm`` package to use this feature."
555+
),
556+
)
527557
def_cell_magic(line,query):
528558
"""Underlying function for bigquery cell magic
529559
@@ -687,12 +717,16 @@ def _cell_magic(line, query):
687717
)
688718
returnquery_job
689719

720+
progress_bar=context.progress_bar_typeorargs.progress_bar_type
721+
690722
ifmax_results:
691723
result=query_job.result(max_results=max_results).to_dataframe(
692-
bqstorage_client=bqstorage_client
724+
bqstorage_client=bqstorage_client,progress_bar_type=progress_bar
693725
)
694726
else:
695-
result=query_job.to_dataframe(bqstorage_client=bqstorage_client)
727+
result=query_job.to_dataframe(
728+
bqstorage_client=bqstorage_client,progress_bar_type=progress_bar
729+
)
696730

697731
ifargs.destination_var:
698732
IPython.get_ipython().push({args.destination_var:result})

‎tests/unit/test_magics.py‎

Lines changed: 69 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -623,7 +623,7 @@ def warning_match(warning):
623623
assertclient_info.user_agent=="ipython-"+IPython.__version__
624624

625625
query_job_mock.to_dataframe.assert_called_once_with(
626-
bqstorage_client=bqstorage_instance_mock
626+
bqstorage_client=bqstorage_instance_mock,progress_bar_type="tqdm"
627627
)
628628

629629
assertisinstance(return_value,pandas.DataFrame)
@@ -665,7 +665,9 @@ def test_bigquery_magic_with_rest_client_requested(monkeypatch):
665665
return_value=ip.run_cell_magic("bigquery","--use_rest_api",sql)
666666

667667
bqstorage_mock.assert_not_called()
668-
query_job_mock.to_dataframe.assert_called_once_with(bqstorage_client=None)
668+
query_job_mock.to_dataframe.assert_called_once_with(
669+
bqstorage_client=None,progress_bar_type="tqdm"
670+
)
669671

670672
assertisinstance(return_value,pandas.DataFrame)
671673

@@ -1167,6 +1169,71 @@ def test_bigquery_magic_w_maximum_bytes_billed_w_context_setter():
11671169
assertsent_config["maximumBytesBilled"]=="10203"
11681170

11691171

1172+
@pytest.mark.usefixtures("ipython_interactive")
1173+
@pytest.mark.skipif(pandasisNone,reason="Requires `pandas`")
1174+
deftest_bigquery_magic_w_progress_bar_type_w_context_setter(monkeypatch):
1175+
ip=IPython.get_ipython()
1176+
ip.extension_manager.load_extension("google.cloud.bigquery")
1177+
magics.context._project=None
1178+
1179+
magics.context.progress_bar_type="tqdm_gui"
1180+
1181+
mock_credentials=mock.create_autospec(
1182+
google.auth.credentials.Credentials,instance=True
1183+
)
1184+
1185+
# Set up the context with monkeypatch so that it's reset for subsequent
1186+
# tests.
1187+
monkeypatch.setattr(magics.context,"_credentials",mock_credentials)
1188+
1189+
# Mock out the BigQuery Storage API.
1190+
bqstorage_mock=mock.create_autospec(bigquery_storage.BigQueryReadClient)
1191+
bqstorage_client_patch=mock.patch(
1192+
"google.cloud.bigquery_storage.BigQueryReadClient",bqstorage_mock
1193+
)
1194+
1195+
sql="SELECT 17 AS num"
1196+
result=pandas.DataFrame([17],columns=["num"])
1197+
run_query_patch=mock.patch(
1198+
"google.cloud.bigquery.magics.magics._run_query",autospec=True
1199+
)
1200+
query_job_mock=mock.create_autospec(
1201+
google.cloud.bigquery.job.QueryJob,instance=True
1202+
)
1203+
query_job_mock.to_dataframe.return_value=result
1204+
withrun_query_patchasrun_query_mock,bqstorage_client_patch:
1205+
run_query_mock.return_value=query_job_mock
1206+
1207+
return_value=ip.run_cell_magic("bigquery","--use_rest_api",sql)
1208+
1209+
bqstorage_mock.assert_not_called()
1210+
query_job_mock.to_dataframe.assert_called_once_with(
1211+
bqstorage_client=None,progress_bar_type=magics.context.progress_bar_type
1212+
)
1213+
1214+
assertisinstance(return_value,pandas.DataFrame)
1215+
1216+
1217+
@pytest.mark.usefixtures("ipython_interactive")
1218+
deftest_bigquery_magic_with_progress_bar_type():
1219+
ip=IPython.get_ipython()
1220+
ip.extension_manager.load_extension("google.cloud.bigquery")
1221+
magics.context.progress_bar_type=None
1222+
1223+
run_query_patch=mock.patch(
1224+
"google.cloud.bigquery.magics.magics._run_query",autospec=True
1225+
)
1226+
withrun_query_patchasrun_query_mock:
1227+
ip.run_cell_magic(
1228+
"bigquery","--progress_bar_type=tqdm_gui","SELECT 17 as num"
1229+
)
1230+
1231+
progress_bar_used=run_query_mock.mock_calls[1][2]["progress_bar_type"]
1232+
assertprogress_bar_used=="tqdm_gui"
1233+
# context progress bar type should not change
1234+
assertmagics.context.progress_bar_typeisNone
1235+
1236+
11701237
@pytest.mark.usefixtures("ipython_interactive")
11711238
deftest_bigquery_magic_with_project():
11721239
ip=IPython.get_ipython()

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp