deepmd/utils/argcheck.py (1)

2308-2340: Constrain types for model_branch_alias and clarify schema

model_branch_alias is declared as a bare list, which weakens validation. Since aliases are strings, prefer list[str] for stricter schema checking. info can remain as dict (free-form metadata), which aligns with the intended flexibility.

Apply this diff:

-            Argument(
-                "model_branch_alias",
-                list,
-                optional=True,
-                default=[],
-                doc=doc_only_pt_supported + doc_model_branch_alias,
-            ),
+            Argument(
+                "model_branch_alias",
+                list[str],
+                optional=True,
+                default=[],
+                doc=doc_only_pt_supported + doc_model_branch_alias,
+            ),

deepmd/utils/model_branch_dict.py (3)

88-96: Avoid dict.keys() in iteration for simplicity

Use direct iteration over the dict to satisfy linters and improve clarity.

Apply this diff:

-        for _, payload in self.data.items():
-            info = payload.get("info") or {}
-            for k in info.keys():
+        for _, payload in self.data.items():
+            info = payload.get("info") or {}
+            for k in info:
                 if k not in seen:
                     seen.add(k)
                     self.info_keys.append(k)

---

166-176: Remove unused loop index in enumerate

The loop index i is unused. Drop enumerate for readability.

Apply this diff:

-        for i, row_cells in enumerate(wrapped_rows):
+        for row_cells in wrapped_rows:
             # Determine the maximum number of wrapped lines in this row
             max_lines = max(len(cell) for cell in row_cells)

---

44-46: Remove meta comment

The comment “generated with GPT for formatted print” is not actionable and can be confusing in source control history.

Apply this diff:

-# generated with GPT for formatted print

deepmd/entrypoints/show.py (1)

40-48: Surface aliases alongside branches in the log

Since selection accepts aliases, include them explicitly in the “Available” message to reduce confusion.

Apply this diff:

-        log.info(
-            f"Available model branches are {model_branches}, "
-            f"where 'RANDOM' means using a randomly initialized fitting net."
-        )
+        aliases = [k for k in model_alias_dict.keys() if k not in model_params["model_dict"]]
+        log.info(
+            f"Available model branches: {model_branches}; aliases: {aliases}. "
+            f"'RANDOM' means using a randomly initialized fitting net."
+        )

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

source/tests/pt/test_dp_show.py (3)

59-61: Fix assertions that accidentally test only the last substring.

Using chained string "and" returns the last string, so the test only checks one token, weakening coverage.

Apply this diff:

-        assert (
-            "{'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut': 4.0"
-        ) in results[2]
+        assert all(
+            s in results[2]
+            for s in ("{'type': 'se_e2_a'", "'sel': [46, 92, 4]", "'rcut': 4.0")
+        )

---

185-195: Fix multitask descriptor assertions to check all expected substrings.

Same issue: chained constants via "and" reduce to the last string.

Apply this diff:

-        assert (
-            "model_1"
-            and "'type': 'se_e2_a'"
-            and "'sel': [46, 92, 4]"
-            and "'rcut_smth': 0.5"
-        ) in results[4]
+        assert all(
+            s in results[4]
+            for s in ("model_1", "'type': 'se_e2_a'", "'sel': [46, 92, 4]", "'rcut_smth': 0.5")
+        )
-        assert (
-            "model_2"
-            and "'type': 'se_e2_a'"
-            and "'sel': [46, 92, 4]"
-            and "'rcut_smth': 0.5"
-        ) in results[5]
+        assert all(
+            s in results[5]
+            for s in ("model_2", "'type': 'se_e2_a'", "'sel': [46, 92, 4]", "'rcut_smth': 0.5")
+        )

---

227-228: Fix singletask (frozen) descriptor assertion to check all substrings.

Same chained "and" issue here.

Apply this diff:

-        assert (
-            "'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut_smth': 0.5"
-        ) in results[2]
+        assert all(
+            s in results[2]
+            for s in ("'type': 'se_e2_a'", "'sel': [46, 92, 4]", "'rcut_smth': 0.5")
+        )

source/tests/pt/test_dp_show.py (2)

169-175: Robustly filter table output in multitask show tests.

Good call filtering ASCII tables and headers to stabilize assertions across richer show outputs.

For reuse and clarity, consider extracting the filter into a small helper to avoid repeating magic substrings:

def _filter_show_output(lines: list[str]) -> list[str]:
    banned = ("DEEPMD WARNING", "|", "+-", "Detailed information")
    return [ln for ln in lines if all(b not in ln for b in banned)]

Then use:

results = _filter_show_output(f.getvalue().split("\n")[:-1])

---

163-175: Optional: add a targeted assertion that the alias table is emitted.

Given the new alias-aware table in show, consider a small focused test that asserts the presence of "Detailed information" and expected alias headers when model_dict contains alias/info, to guard against regressions in the table path. Keep the existing filtered test for the non-table content.

I can draft a self-contained test method that runs dp --pt show with model-branch and asserts the table header and a couple of alias values without depending on exact column widths. Want me to propose it?

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these settings in your CodeRabbit configuration.

deepmd/entrypoints/show.py (1)

40-40: Past suggestion addressed: unpack from get_model_dict correctly

Using "_, model_branch_dict = get_model_dict(...)" matches the earlier review suggestion. LGTM.

deepmd/entrypoints/show.py (4)

31-49: Nice integration of alias-aware branch listing with a detailed table

• Properly guards the "model-branch" path to multitask models.
• The human-readable list plus the ASCII table is a good UX improvement.

One nit: you list "RANDOM" in the short list but it’s absent from the table. Consider adding a row for clarity (see next comment with a concrete diff).

---

45-48: Optionally include a 'RANDOM' row in the table for consistency

Right now "RANDOM" appears in the plain list but not in the table, which may confuse users. You can inject a minimal placeholder row before rendering.

Apply this diff:

         _, model_branch_dict = get_model_dict(model_params["model_dict"])
         log.info(
             f"Available model branches are {model_branches}, "
             f"where 'RANDOM' means using a randomly initialized fitting net."
         )
-        log.info(
+        # Optionally reflect RANDOM in the table as a placeholder
+        model_branch_dict = dict(model_branch_dict)
+        model_branch_dict.setdefault(
+            "RANDOM",
+            {"alias": [], "info": {"description": "Randomly initialized fitting net"}},
+        )
+        log.info(
             "Detailed information: \n"
             + OrderedDictTableWrapper(model_branch_dict).as_table()
         )

---

91-106: Validate and normalize info.observed_type; dedupe to avoid double-counting

If info.observed_type is mis-typed (e.g., a string) or contains duplicates, logs and totals become misleading. Light validation and normalization would make this robust while preserving order.

Apply this diff:

-                if (
-                    model_params["model_dict"][branch]
-                    .get("info", {})
-                    .get("observed_type", None)
-                    is not None
-                ):
-                    observed_type_list = model_params["model_dict"][branch]["info"][
-                        "observed_type"
-                    ]
-                    observed_types = {
-                        "type_num": len(observed_type_list),
-                        "observed_type": observed_type_list,
-                    }
+                if (
+                    model_params["model_dict"][branch]
+                    .get("info", {})
+                    .get("observed_type", None)
+                    is not None
+                ):
+                    observed_type_list = model_params["model_dict"][branch]["info"]["observed_type"]
+                    # Normalize: accept a single string, preserve order, remove dups
+                    if isinstance(observed_type_list, (str, bytes)):
+                        observed_type_list = [observed_type_list]
+                    if not isinstance(observed_type_list, (list, tuple)):
+                        log.warning(
+                            "Branch %s: info.observed_type must be a list of strings; falling back to model introspection.",
+                            branch,
+                        )
+                        tmp_model = DeepEval(INPUT, head=branch, no_jit=True)
+                        observed_types = tmp_model.get_observed_types()
+                    else:
+                        unique_observed = list(dict.fromkeys(observed_type_list))
+                        observed_types = {
+                            "type_num": len(unique_observed),
+                            "observed_type": unique_observed,
+                        }
                 else:
                     tmp_model = DeepEval(INPUT, head=branch, no_jit=True)
                     observed_types = tmp_model.get_observed_types()

---

105-106: Avoid repeated model instantiation when falling back to introspection

Creating a new DeepEval per branch can be expensive. If the API allows, consider reusing a single instance and switching heads, or caching results per branch to cut load time.

If you want, I can sketch a small cache layer keyed by branch to avoid redundant DeepEval(...) calls when many branches miss observed_type in info.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

deepmd/utils/model_branch_dict.py (1)

32-41: Replace assert with explicit exception and harden None handling for aliases/info

Asserts can be stripped with Python -O, turning user input errors into silent misbehavior. Also, model_branch_alias or info may be None, which would raise at iteration time.

Apply this diff:

-        alias_list = model_dict[key].get("model_branch_alias", [])
+        alias_list = model_dict[key].get("model_branch_alias") or []
         model_branch_dict[key]["alias"] = alias_list
-        branch_info = model_dict[key].get("info", {})
+        branch_info = model_dict[key].get("info") or {}
         model_branch_dict[key]["info"] = branch_info
         for alias in alias_list:
-            assert alias not in model_alias_dict, (
-                f"Alias {alias} for model_branch {key} already exists in model_branch {model_alias_dict[alias]}!"
-            )
+            if alias in model_alias_dict:
+                raise ValueError(
+                    f"Alias '{alias}' for model_branch '{key}' already exists in model_branch "
+                    f"'{model_alias_dict[alias]}'!"
+                )
             model_alias_dict[alias] = key

deepmd/utils/model_branch_dict.py (6)

91-97: Drop redundant .keys() usage when iterating and checking membership

Minor cleanup to satisfy linters (SIM118) and improve readability.

Apply this diff:

-        for _, payload in self.data.items():
-            info = payload.get("info") or {}
-            for k in info.keys():
-                if k not in seen:
-                    seen.add(k)
-                    self.info_keys.append(k)
+        for _, payload in self.data.items():
+            info = payload.get("info") or {}
+            for k in info:
+                if k not in seen:
+                    seen.add(k)
+                    self.info_keys.append(k)

---

127-131: Use direct dict iteration to compute branch width

Matches Ruff’s SIM118 and is a minor readability win.

Apply this diff:

-        for branch in self.data.keys():
+        for branch in self.data:
             branch_col_width = max(branch_col_width, len(str(branch)))

---

133-137: Alias column may wrap unexpectedly; compute width from the rendered alias string

You calculate width from the maximum single alias, but display a comma-joined string. For multiple aliases, the actual cell content can exceed the computed width, causing unintended wrapping.

Apply this diff:

-        alias_col_width = len(self.headers[1])  # "Alias"
-        for payload in self.data.values():
-            alias_list = payload.get("alias", [])
-            for alias in alias_list:
-                alias_col_width = max(alias_col_width, len(str(alias)))
+        alias_col_width = len(self.headers[1])  # "Alias"
+        for payload in self.data.values():
+            alias_str = ", ".join(map(str, payload.get("alias", [])))
+            alias_col_width = max(alias_col_width, len(alias_str))

Also applies to: 145-147

---

198-198: Rename unused loop index

The loop variable i is unused (B007). Rename to _ to signal intentional discard.

Apply this diff:

-        for i, row_cells in enumerate(wrapped_rows):
+        for _, row_cells in enumerate(wrapped_rows):

---

98-101: Nit: Header text mismatch with docstring

Docstring says “Model Branch Name” while header uses “Model Branch”. Pick one for consistency.

Example:

-        self.headers: list[str] = ["Model Branch", "Alias", *self.info_keys]
+        self.headers: list[str] = ["Model Branch Name", "Alias", *self.info_keys]

---

12-26: Add type hints for clarity and downstream tooling

get_model_dict lacks type hints; adding them improves IDE help and self-documentation.

Apply this diff:

-from typing import (
-    Any,
-    Optional,
-)
+from typing import Any, Optional, Dict, Tuple, OrderedDict as TOrderedDict
@@
-def get_model_dict(model_dict):
+def get_model_dict(
+    model_dict: Dict[str, Dict[str, Any]]
+) -> Tuple[Dict[str, str], Dict[str, Dict[str, Any]]]:

deepmd/pt/infer/deep_eval.py (3)

125-127: Drop unused variable to avoid linter noise

model_branch_dict is not used in this method.

Apply this diff:

-                model_alias_dict, model_branch_dict = get_model_dict(
+                model_alias_dict, _ = get_model_dict(
                     self.input_param["model_dict"]
                 )

---

140-148: Case-insensitive fallback should only run for string inputs

If head is not a string, calling .lower() will raise. Also, keep behavior unchanged for non-strings.

Apply this diff:

-                if head not in model_alias_dict:
-                    # preprocess with potentially case-insensitive input
-                    head_lower = head.lower()
-                    for mk in model_alias_dict:
-                        if mk.lower() == head_lower:
-                            # mapped the first matched head
-                            head = mk
-                            break
+                if isinstance(head, str) and head not in model_alias_dict:
+                    # case-insensitive fallback: map to the first matching key
+                    head_lower = head.lower()
+                    for mk in model_alias_dict:
+                        if mk.lower() == head_lower:
+                            head = mk
+                            break

---

153-163: Optional: record the resolved branch and log selection

Expose the resolved branch for downstream consumers (debugging, telemetry, or UIs) and log it once.

Apply this diff:

-                head = model_alias_dict[head]
+                head = model_alias_dict[head]
+                self.selected_model_branch = head  # for introspection
+                log.info(f"Selected multitask head: {head}")

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.