Fix Arg[type] subscript type not captured for Arrow type inference

The Arg[str] subscript type was not being captured at runtime, causing
argument types to show as "null" in catalog function metadata.

Changes:
- Add _ArgFactory class that captures the type parameter when Arg[type]
  is used and creates Arg instances with _type_param set
- Add _type_param slot to Arg class to store the captured type
- Update extract_argument_specs to use _type_param when determining
  Arrow type (priority: explicit arrow_type > _type_param > type hint)
- Update test to use a custom type not in PYTHON_TO_ARROW

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: rustyconover

tests/test_argument_spec.py

@@ -490,10 +490,14 @@ class TestExtractArgumentSpecsValidation:
     def test_missing_type_hint_warns(self) -> None:
         """Missing type hint and no arrow_type should issue a warning."""
 
+        # Use a custom class not in PYTHON_TO_ARROW to trigger warning
+        class CustomType:
+            pass
+
         class FunctionWithArg(TableInOutFunction):
-            count = Arg[int](0)  # No type annotation, no arrow_type
+            count = Arg[CustomType](0)  # Type not in PYTHON_TO_ARROW
 
-        # Should warn about missing type
+        # Should warn about missing type (CustomType is not in PYTHON_TO_ARROW)
         with pytest.warns(UserWarning, match="Cannot determine Arrow type"):
             specs = extract_argument_specs(FunctionWithArg)
 

vgi/argument_spec.py

@@ -344,13 +344,21 @@ class MyFunction(TableInOutFunction):
 
                 # Determine Arrow type using priority order:
                 # 1. Explicit arrow_type on Arg
-                # 2. Type hint with PYTHON_TO_ARROW mapping
-                # 3. Default to pa.null() with warning
+                # 2. Type parameter from Arg[type] subscript (e.g., Arg[str])
+                # 3. Type hint with PYTHON_TO_ARROW mapping
+                # 4. Default to pa.null() with warning
                 arrow_type: pa.DataType
                 if arg.arrow_type is not None:
                     arrow_type = arg.arrow_type
                 elif is_table_input or is_any_type:
                     arrow_type = pa.null()
+                elif (
+                    hasattr(arg, "_type_param")
+                    and arg._type_param is not None
+                    and arg._type_param in PYTHON_TO_ARROW
+                ):
+                    # Use type from Arg[type] subscript
+                    arrow_type = PYTHON_TO_ARROW[arg._type_param]
                 elif hint is not None and hint in PYTHON_TO_ARROW:
                     arrow_type = PYTHON_TO_ARROW[hint]
                 else:

vgi/arguments.py

@@ -486,6 +486,73 @@ def _suggest_similar_choices(self) -> list[Any]:
 ArgT = TypeVar("ArgT")
 
 
+class _ArgFactory:
+    """Factory returned by Arg[type] to capture the type parameter.
+
+    This allows Arg[str](0) to create an Arg instance with _type_param=str,
+    which can be used by extract_argument_specs to infer the Arrow type.
+    """
+
+    __slots__ = ("_type_param",)
+
+    def __init__(self, type_param: type) -> None:
+        self._type_param = type_param
+
+    def __call__(
+        self,
+        position: int | str,
+        *,
+        default: Any = _MISSING,
+        doc: str = "",
+        ge: float | int | None = None,
+        le: float | int | None = None,
+        gt: float | int | None = None,
+        lt: float | int | None = None,
+        choices: Sequence[Any] | None = None,
+        pattern: str | None = None,
+        varargs: bool = False,
+        arrow_type: pa.DataType | None = None,
+    ) -> "Arg[Any]":
+        """Create an Arg instance with the captured type parameter."""
+        arg: Arg[Any] = Arg.__new__(Arg)
+        # Manually call __init__ logic since we're using __new__
+        # Validate constraint combinations
+        if ge is not None and gt is not None:
+            raise ValueError("Cannot specify both 'ge' and 'gt'")
+        if le is not None and lt is not None:
+            raise ValueError("Cannot specify both 'le' and 'lt'")
+        if varargs:
+            if isinstance(position, str):
+                raise ValueError(
+                    "varargs=True requires a positional argument (int), not named"
+                )
+            if default is not _MISSING:
+                raise ValueError(
+                    "varargs=True cannot have a default value "
+                    "(requires at least 1 value)"
+                )
+
+        arg.position = position
+        arg.default = default
+        arg.doc = doc
+        arg.ge = ge
+        arg.le = le
+        arg.gt = gt
+        arg.lt = lt
+        arg.choices = choices
+        arg.pattern = pattern
+        arg.varargs = varargs
+        arg.arrow_type = arrow_type
+        arg._name = None
+        arg._compiled_pattern = None
+        arg._type_param = self._type_param
+
+        if pattern is not None:
+            arg._compiled_pattern = re.compile(pattern)
+
+        return arg
+
+
 class Arg[ArgT]:
     """Descriptor for declarative argument parsing with optional validation.
 
@@ -559,6 +626,7 @@ def transform(self, batch):
         "arrow_type",
         "_name",
         "_compiled_pattern",
+        "_type_param",
     )
 
     def __init__(
@@ -629,11 +697,21 @@ def __init__(
         self.arrow_type = arrow_type
         self._name: str | None = None
         self._compiled_pattern: re.Pattern[str] | None = None
+        self._type_param: type | None = None
 
         # Pre-compile pattern for efficiency
         if pattern is not None:
             self._compiled_pattern = re.compile(pattern)
 
+    def __class_getitem__(cls, item: type) -> "_ArgFactory":
+        """Support Arg[type] syntax to capture the type parameter at runtime.
+
+        When you write Arg[str](0), this method is called first with item=str,
+        and returns an _ArgFactory that will create Arg instances with
+        _type_param set to str.
+        """
+        return _ArgFactory(item)
+
     def __set_name__(self, owner: type, name: str) -> None:
         """Store the attribute name when assigned to a class."""
         self._name = name