enh(engines/utils): use tuple/dict comprehensions in PrepareBatchExtraInput

[Zeesejo]

Fixes #8806

Description

This implements the suggestion made by @ericspod in the review of #8747.

In PrepareBatchExtraInput.__call__ (monai/engines/utils.py), replace the imperative list.append loop and dict.update loop with a direct tuple generator expression and dict comprehension, respectively. Also adds explicit type annotations (args_: tuple and kwargs_: dict) for clarity.

Before:

args_ = []
kwargs_ = {}
...
if isinstance(self.extra_keys, (str, list, tuple)):
    for k in ensure_tuple(self.extra_keys):
        args_.append(_get_data(k))
elif isinstance(self.extra_keys, dict):
    for k, v in self.extra_keys.items():
        kwargs_.update({k: _get_data(v)})
return ..., tuple(args_), kwargs_

After:

args_: tuple = ()
kwargs_: dict = {}
...
if isinstance(self.extra_keys, (str, list, tuple)):
    args_ = tuple(_get_data(k) for k in ensure_tuple(self.extra_keys))
elif isinstance(self.extra_keys, dict):
    kwargs_ = {k: _get_data(v) for k, v in self.extra_keys.items()}
return ..., args_, kwargs_

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PrepareBatchExtraInput.__call__ method in monai/engines/utils.py was refactored to use immutable data structures. The args_ variable transitions from a list with successive append calls to a tuple constructed via generator expression. Similarly, kwargs_ now uses a dictionary comprehension for a single assignment instead of incremental update calls. The return statement directly uses the new tuple, eliminating the final conversion step.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title accurately describes the main change: refactoring PrepareBatchExtraInput to use tuple/dict comprehensions instead of imperative loops.
Description check	✅ Passed	Description is complete with context, before/after code examples, and all required checklist items addressed. Properly references linked issue #8806.
Linked Issues check	✅ Passed	PR fully implements the enhancement suggested in #8747 and tracked by #8806: replacing imperative loops with comprehensions and adding type annotations.
Out of Scope Changes check	✅ Passed	Changes are tightly scoped to PrepareBatchExtraInput.call in monai/engines/utils.py with no unrelated modifications.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[benediktjohannes]

@coderabbitai review

[coderabbitai]

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

monai/engines/utils.py (1)

222-223: Tighten the type annotations (optional).

tuple and dict are valid but uninformative. Consider tuple[torch.Tensor, ...] and dict[str, torch.Tensor] to match the return signature on line 215 and what _get_data yields.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@monai/engines/utils.py` around lines 222 - 223, The current annotations for
args_ and kwargs_ are too broad; narrow them to match what _get_data returns and
the function's declared return (see the return signature around line 215).
Change args_: tuple = () to something like args_: tuple[torch.Tensor, ...] = ()
and kwargs_: dict = {} to kwargs_: dict[str, torch.Tensor] = {} (or other more
specific key/value types used by _get_data) so the types reflect actual tensors
and string keys used by the code.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@monai/engines/utils.py`:
- Around line 236-240: The indentation around the branch that handles
self.extra_keys is wrong: fix the over-indented lines so the elif
isinstance(self.extra_keys, dict): and its kwargs_ = {k: _get_data(v) ...} are
aligned with the other conditional branches and move the return statement
(return cast(torch.Tensor, image), cast(torch.Tensor, label), args_, kwargs_)
out of the elif-block to the outer scope of the function so it always executes;
ensure symbols referenced include self.extra_keys, _get_data, kwargs_, args_,
image and label so the dict, list/str/tuple and default branches all return
consistently.

---

Nitpick comments:
In `@monai/engines/utils.py`:
- Around line 222-223: The current annotations for args_ and kwargs_ are too
broad; narrow them to match what _get_data returns and the function's declared
return (see the return signature around line 215). Change args_: tuple = () to
something like args_: tuple[torch.Tensor, ...] = () and kwargs_: dict = {} to
kwargs_: dict[str, torch.Tensor] = {} (or other more specific key/value types
used by _get_data) so the types reflect actual tensors and string keys used by
the code.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 52809bac-7f01-4e6d-9d51-a2e6b88f7193

📥 Commits

Reviewing files that changed from the base of the PR and between 1dc47a5 and 85a1039.

📒 Files selected for processing (1)

• monai/engines/utils.py

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

Broken indentation — file fails to import.

Lines 237 and 240 are over-indented, triggering IndentationError: unindent does not match any outer indentation level (confirmed by Ruff and the premerge/premerge-min pipelines). Beyond the syntax error, the return at line 240 sits at the elif block's indentation, so once parsed it would only execute when extra_keys is a dict — the str/list/tuple branch would fall through with no return.

🔧 Proposed fix

         if isinstance(self.extra_keys, (str, list, tuple)):
             args_ = tuple(_get_data(k) for k in ensure_tuple(self.extra_keys))
-
         elif isinstance(self.extra_keys, dict):
-                        kwargs_ = {k: _get_data(v) for k, v in self.extra_keys.items()}
-
+            kwargs_ = {k: _get_data(v) for k, v in self.extra_keys.items()}

-                return cast(torch.Tensor, image), cast(torch.Tensor, label), args_, kwargs_
+        return cast(torch.Tensor, image), cast(torch.Tensor, label), args_, kwargs_

🧰 Tools

🪛 GitHub Actions: premerge

[error] 240-240: Python import failed with IndentationError: 'unindent does not match any outer indentation level' at line 240.

🪛 GitHub Actions: premerge-min

[error] 240-240: Python failed with IndentationError: "unindent does not match any outer indentation level" during import. Step: "python -c "import monai; monai.config.print_config()"".

🪛 Ruff (0.15.10)

[warning] 240-240: unindent does not match any outer indentation level

(invalid-syntax)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@monai/engines/utils.py` around lines 236 - 240, The indentation around the
branch that handles self.extra_keys is wrong: fix the over-indented lines so the
elif isinstance(self.extra_keys, dict): and its kwargs_ = {k: _get_data(v) ...}
are aligned with the other conditional branches and move the return statement
(return cast(torch.Tensor, image), cast(torch.Tensor, label), args_, kwargs_)
out of the elif-block to the outer scope of the function so it always executes;
ensure symbols referenced include self.extra_keys, _get_data, kwargs_, args_,
image and label so the dict, list/str/tuple and default branches all return
consistently.