fix: reject non-canonical integer keys in tree_unflatten

[Solaris-star]

Proposed changes

tree_unflatten treated all integer-like keys as list indices via int(idx). Keys such as "01" and "1" both parse as 1, which silently shifted later values and inserted a phantom {} (issue #3877).

This PR rejects non-canonical integer keys (str(i) != k) and duplicate indices by raising ValueError, so the existing except ValueError path falls back to a dict tree and preserves values.

Fixes #3877

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[0lekz]

LGTM. Change makes tree_unflatten reject ambiguous keys instead of silently producing a corrupted list, and the regression test covers both the edge case and the canonical behavior. I don't see any compatibility concerns since tree_flatten already emits canonical integer keys

[Solaris-star]

Thanks @0lekz for the review and approval — much appreciated.

Happy to address any further feedback. If this looks good from a maintainer perspective, a merge when convenient would be great.

[Pablosinyores]

Confirmed the bug and the fix behaves as described. Checked out the branch and ran test_tree, all 5 pass, and the collision case now falls back to a dict instead of the shifted [{}, 'a', 'b', 'c'].

Three things worth looking at before this goes in.

black will fail the lint job. CI has not been approved to run on this fork PR yet, so it is not visible yet, but black --check python/tests/test_tree.py reformats it:

-
     def test_tree_unflatten_integer_key_collision(self):
...
+
 if __name__ == "__main__":

One blank line too many before the new method and one too few before if __name__. Running black python/tests/test_tree.py fixes both.

The two ValueError messages can never reach a user. They are raised inside the try whose except ValueError is the existing fall-back-to-dict path, so they are immediately swallowed and the caller just gets a dict. That is a reasonable outcome, and it matches how the function already treats non-integer keys, but it means the carefully worded "collides with index" and "Duplicate integer key" text is unreachable. Either drop the messages and use a bare raise ValueError with a comment saying the fall-back is deliberate, or, if you want the user to actually see the diagnosis, raise outside the try.

Worth flagging that the issue you filed leaned toward raising rather than falling back. Those are genuinely different contracts for someone loading an externally produced checkpoint: raising tells them their keys are malformed, falling back hands them a dict where they expected a list and lets it fail somewhere further away. That is a maintainer call, but it would be good to have it made explicitly rather than inherited from where the raise happens to sit.

The duplicate-index check looks unreachable. children is a dict, so its keys are already distinct strings. If two distinct keys map to the same i, at least one of them is not str(i), so the canonical check raises first. If both satisfy str(i) == k they are the same string and cannot both be present. So seen_indices never fires and could come out, which would make the intent of the remaining check clearer.

Nice find on the root cause, the range(i - len(l)) analysis in the issue is exactly right.