使tqdm等函数在Python中接受不同类型的参数

Question

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:

Answer 1

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!