kkgarg
kkgarg

Reputation: 1376

Making tqdm and the like functions accept different types of parameters in Python

This is not specific to tqdm but a generic question about passing parameters to a function in Python. I want to achieve the following functionality without having to make copies of the entire-block under tqdm. Any help will be greatly appreciated.

if flag == True:
    with tqdm(dataloader, total=args.num_train_batches) as pbar:
else:
    with tqdm(dataloader) as pbar:

More specifically, can I pass parameters in a way like this?

if flag == True:
    tqdm_args = dataloader, total=args.num_train_batches
else:
    tqdm_args = dataloader
with tqdm(tqdm_args) as pbar:

Upvotes: 2

Views: 468

Answers (1)

JaonHax
JaonHax

Reputation: 346

This is actually fairly simple to do, as it seems they thought of this when making Python. You can use Python's ternary operator to do this, condensing what you have above to a single line:

with tqdm(dataloader, total=args.num_train_batches if flag else None) as pbar:
  # ...

Edit: to answer with your preferred method you mentioned, yes. That's also possible. If you put those arguments into a list (or dictionary, if you have keyword args) and then put a * (or ** for a dictionary) in front of the list's name when calling the function, it unpacks the list into a set of arguments.

Example using a list:

if flag: # if flag is a boolean, putting "== True" does nothing
    tqdm_args = [dataloader, None, args.num_train_batches]
else:
    tqdm_args = [dataloader]
with tqdm(*tqdm_args) as pbar:
    # ...

Example with a dictionary:

if flag:
  tqdm_kwargs = {"iterable": dataloader, "total": args.num_train_batches}
else:
  tqdm_kwargs = {"iterable": dataloader}
with tqdm(**tqdm_kwargs) as pbar:
  # ...

Happy to be of assistance!

Upvotes: 4

Related Questions