Skip to content

Commit

Permalink
Mute progress bars with MUTE_PROGRESS_BARS env variable (#545)
Browse files Browse the repository at this point in the history
* Enable progress bar mute with env variable

* Temp remove windows test

* Temp remove macos testing
  • Loading branch information
munrojm committed Mar 4, 2022
1 parent d64c402 commit bc3094c
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 8 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/testing.yml
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ jobs:
strategy:
max-parallel: 2
matrix:
os: [ubuntu-latest, windows-latest, macos-latest]
os: [ubuntu-latest]
python-version: [3.8, 3.9]

runs-on: ${{ matrix.os }}
Expand Down
23 changes: 16 additions & 7 deletions src/mp_api/core/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -476,9 +476,13 @@ def _submit_requests(
num_docs_needed = min((max_pages * chunk_size), total_num_docs)

# Setup progress bar
pbar = tqdm(
desc=f"Retrieving {self.document_model.__name__} documents", # type: ignore
total=num_docs_needed,
pbar = (
tqdm(
desc=f"Retrieving {self.document_model.__name__} documents", # type: ignore
total=num_docs_needed,
)
if not MAPIClientSettings().MUTE_PROGRESS_BARS
else None
)

initial_data_length = len(total_data["data"])
Expand All @@ -487,15 +491,18 @@ def _submit_requests(
if initial_data_length >= num_docs_needed or num_chunks == 1:
new_total_data = copy(total_data)
new_total_data["data"] = total_data["data"][:num_docs_needed]
pbar.update(num_docs_needed)
pbar.close()

if pbar is not None:
pbar.update(num_docs_needed)
pbar.close()
return new_total_data

# otherwise, prepare to paginate in parallel
if chunk_size is None:
raise ValueError("A chunk size must be provided to enable pagination")

pbar.update(initial_data_length)
if pbar is not None:
pbar.update(initial_data_length)

# Warning to select specific fields only for many results
if criteria.get("all_fields", False) and (total_num_docs / chunk_size > 10):
Expand Down Expand Up @@ -546,7 +553,8 @@ def _submit_requests(
if "meta" in data:
total_data["meta"]["time_stamp"] = data["meta"]["time_stamp"]

pbar.close()
if pbar is not None:
pbar.close()

return total_data

Expand Down Expand Up @@ -780,6 +788,7 @@ def search(
A generic search method to retrieve documents matching specific parameters.
Arguments:
mute (bool): Whether to mute progress bars.
num_chunks (int): Maximum number of chunks of data to yield. None will yield all possible.
chunk_size (int): Number of data entries per chunk.
all_fields (bool): Set to False to only return specific fields of interest. This will
Expand Down
4 changes: 4 additions & 0 deletions src/mp_api/core/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,3 +44,7 @@ class MAPIClientSettings(BaseSettings):
)

MAX_RETRIES: int = Field(3, description="Maximum number of retries for requests.")

MUTE_PROGRESS_BARS: bool = Field(
False, description="Whether to mute progress bars when data is retrieved.",
)

0 comments on commit bc3094c

Please sign in to comment.