fix(arguments): satisfy strict mypy + ruff for union-tag decode

CI lint/type gates (not run locally before the prior push) flagged the union
decode: drop the redundant UnionType cast, coerce type_code to int (the pyarrow
stub types it as str though it is an int at runtime) so the type_codes lookup
type-checks, and add docstrings to the new union tests.

ruff, mypy (vgi/), and the union tests all clean locally.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: rustyconover

tests/test_union_argument.py

@@ -17,8 +17,8 @@
 _SVC = pa.struct([pa.field("C", pa.list_(pa.float64())), pa.field("kernel", pa.list_(pa.string()))])
 
 
-def _union_scalar(tag: str, code: int, members: dict) -> pa.UnionScalar:
-    """Build a one-element sparse-union scalar with named members, active = ``tag``."""
+def _union_scalar(code: int, members: dict) -> pa.UnionScalar:
+    """Build a one-element sparse-union scalar with named members, active = ``code``."""
     arr = pa.UnionArray.from_sparse(
         pa.array([code], type=pa.int8()),
         [
@@ -32,23 +32,24 @@ def _union_scalar(tag: str, code: int, members: dict) -> pa.UnionScalar:
 
 
 def test_union_arg_preserves_tag() -> None:
-    scalar = _union_scalar(
-        "random_forest_classifier", 0, {"random_forest_classifier": {"n_estimators": [100, 300], "max_depth": [3, 5]}}
-    )
+    """A union argument decodes to a TaggedUnion carrying the active member name."""
+    scalar = _union_scalar(0, {"random_forest_classifier": {"n_estimators": [100, 300], "max_depth": [3, 5]}})
     got = Arguments(named={"config": scalar}).get("config")
     assert isinstance(got, TaggedUnion)
     assert got.tag == "random_forest_classifier"
     assert got.value == {"n_estimators": [100, 300], "max_depth": [3, 5]}
 
 
 def test_union_arg_other_member() -> None:
-    scalar = _union_scalar("svc", 1, {"svc": {"C": [1.0, 10.0], "kernel": ["rbf", "linear"]}})
+    """The tag reflects whichever union member is set."""
+    scalar = _union_scalar(1, {"svc": {"C": [1.0, 10.0], "kernel": ["rbf", "linear"]}})
     got = Arguments(named={"config": scalar}).get("config")
     assert got.tag == "svc"
     assert got.value == {"C": [1.0, 10.0], "kernel": ["rbf", "linear"]}
 
 
 def test_non_union_args_unchanged() -> None:
+    """Non-union arguments still decode via plain as_py()."""
     args = Arguments(named={"n": pa.scalar(5), "s": pa.scalar("hi")}, positional=(pa.scalar(1.5),))
     assert args.get("n") == 5
     assert args.get("s") == "hi"

vgi/arguments.py

@@ -194,23 +194,28 @@ class TaggedUnion:
     value: Any
 
 
-def _scalar_to_py(scalar: "Scalar") -> Any:
+def _scalar_to_py(scalar: "Scalar[Any]") -> Any:
     """Convert an argument scalar to a Python value, preserving union tags.
 
     Identical to ``scalar.as_py()`` for every type except unions: a
     ``UnionScalar`` is decoded to a [`TaggedUnion`][] so the member
     discriminator (which ``as_py()`` discards) is retained.
+
+    Args:
+        scalar: The argument scalar to convert.
+
+    Returns:
+        ``scalar.as_py()`` for non-union scalars; a [`TaggedUnion`][] for unions.
+
     """
     if isinstance(scalar, pa.UnionScalar):
+        # Map the active ``type_code`` to its member field name via the union
+        # type's parallel ``type_codes`` / ``field()``. (``type_code`` is coerced
+        # to int — it is an integer at runtime regardless of the stub's typing.)
         union_type = scalar.type
-        tag = next(
-            (
-                union_type.field(i).name
-                for i in range(union_type.num_fields)
-                if union_type.type_codes[i] == scalar.type_code
-            ),
-            None,
-        )
+        type_codes = list(union_type.type_codes)
+        code = int(scalar.type_code)
+        tag = union_type.field(type_codes.index(code)).name if code in type_codes else None
         inner = scalar.value
         return TaggedUnion(tag=tag, value=inner.as_py() if inner is not None else None)
     return scalar.as_py()