refactor(extensions)!: centralize argument extraction

API.md

@@ -101,7 +101,7 @@ Substrait has two extension mechanisms:
 
 - **Simple extensions** declare extension functions, types, and type variations.
   `substrait-explain` reads and writes these declarations in the `===
-  Extensions` section and uses them to resolve anchors while parsing and
+Extensions` section and uses them to resolve anchors while parsing and
   formatting expressions and types. No YAML files are read; substrait-explain
   relies on the protobufs / text format itself for function names, and does not
   validate type signatures exist / match extensions.
@@ -137,9 +137,10 @@ extension arguments. `substrait-explain` handles parsing and rendering the text
 syntax around those arguments:
 
 - `name()` - The extension name used in text (e.g., `"ParquetScan"`)
-- `from_args(args)` - Parse text arguments into your type
-- `to_args(&self, context)` - Convert your type to text arguments, optionally
-  using information about relation inputs.
+- `from_args(args)` - Build your extension value from text arguments with
+  `ArgsAccess`
+- `to_args(&self, context)` - Convert your extension value to text arguments,
+  optionally using information about relation inputs
 
 The extension API works across three representations:
 
@@ -151,7 +152,7 @@ The extension API works across three representations:
 Untyped scalar extension literals such as `2` or `'path'` are represented as
 scalar `ExtensionValue` variants and render without expression type suffixes,
 even in verbose output. The same values can still be requested as `Expr` through
-`ArgsExtractor`, which widens them to default non-nullable Substrait literal
+`ArgsAccess`, which widens them to default non-nullable Substrait literal
 expressions. Typed literals, field references, function calls, and casts are
 represented as expression values.
 
@@ -164,11 +165,21 @@ The `to_args` method receives an `ExtensionContext`. Relation extensions get one
 emitted column count, including any output mapping applied by that child. Other
 extension namespaces receive an empty input slice.
 
-Use `ArgsExtractor` for convenient argument parsing:
+### Reading Extension Arguments
 
-- `extractor.expect_named::<T>(name)` - Required argument
-- `extractor.get_named::<T>(name)?` - Optional argument, returning `Option<T>`
-- `extractor.check_exhausted()` - Verify no unexpected arguments
+The registry passes [`ArgsAccess`] to your [`Explainable::from_args`]
+implementation. It provides access to positional arguments, named arguments,
+and the relation's output columns declared by custom relations.
+
+If `from_args` returns successfully, the registry rejects any positional or
+named arguments the implementation did not access. If `from_args` returns an
+error, the registry preserves that error without checking for unhandled
+arguments. Output columns are relation metadata and are not included in this
+check.
+
+Call [`ExtensionArgs::parse`](extensions::ExtensionArgs::parse) to perform the
+same checked conversion directly, such as when testing an `Explainable`
+implementation.
 
 ### Extension Namespaces
 
@@ -189,7 +200,8 @@ Register extensions to the appropriate namespace:
 ```rust,no_run
 # use prost::{Message, Name};
 # use substrait_explain::extensions::{
-#     Explainable, ExtensionArgs, ExtensionContext, ExtensionError, ExtensionRegistry,
+#     ArgsAccess, Explainable, ExtensionArgs, ExtensionContext, ExtensionError,
+#     ExtensionRegistry,
 # };
 #[derive(Clone, PartialEq, Message)]
 pub struct MySourceConfig {
@@ -204,7 +216,9 @@ impl Name for MySourceConfig {
 }
 # impl Explainable for MySourceConfig {
 #     fn name() -> &'static str { "MySource" }
-#     fn from_args(_: &ExtensionArgs) -> Result<Self, ExtensionError> { Ok(Self::default()) }
+#     fn from_args(_: &mut ArgsAccess<'_>) -> Result<Self, ExtensionError> {
+#         Ok(Self::default())
+#     }
 #     fn to_args(
 #         &self,
 #         _context: &ExtensionContext<'_>,

examples/extensions.rs

@@ -15,8 +15,8 @@ use prost::{Message, Name};
 use substrait::proto::{self, Plan, PlanRel, Rel, plan_rel, rel};
 use substrait_explain::extensions::any::AnyRef;
 use substrait_explain::extensions::{
-    AnyConvertible, Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError,
-    ExtensionRegistry,
+    AnyConvertible, ArgsAccess, Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext,
+    ExtensionError, ExtensionRegistry,
 };
 use substrait_explain::{OutputOptions, Parser, format_with_registry};
 
@@ -69,19 +69,15 @@ impl Explainable for ParquetScanConfig {
         "TypedParquetScan"
     }
 
-    fn from_args(args: &ExtensionArgs) -> Result<Self, ExtensionError> {
-        let mut extractor = args.extractor();
+    fn from_args(args: &mut ArgsAccess<'_>) -> Result<Self, ExtensionError> {
         // path is required
-        let path: &str = extractor.expect_named("path")?;
+        let path: &str = args.expect_named("path")?;
         // batch_size and use_dictionary are optional, with default values
-        let batch_size: i64 = extractor.get_named("batch_size")?.unwrap_or(1024);
-        let use_dictionary: bool = extractor.get_named("use_dictionary")?.unwrap_or(true);
-
-        // Validate there are no other named arguments
-        extractor.check_exhausted()?;
+        let batch_size: i64 = args.get_named("batch_size")?.unwrap_or(1024);
+        let use_dictionary: bool = args.get_named("use_dictionary")?.unwrap_or(true);
 
         let selected_columns = args
-            .output_columns
+            .output_columns()
             .iter()
             .map(|column| match column {
                 ExtensionColumn::Named { name, r#type } => Ok(ParquetColumn {

src/cli.rs

@@ -480,7 +480,7 @@ mod tests {
 
     use super::*;
     use crate::extensions::{
-        Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError,
+        ArgsAccess, Explainable, ExtensionArgs, ExtensionColumn, ExtensionContext, ExtensionError,
     };
     use crate::fixtures::parse_type;
     use crate::parse;
@@ -907,10 +907,8 @@ Root[result]
             "TestSource"
         }
 
-        fn from_args(args: &ExtensionArgs) -> Result<Self, ExtensionError> {
-            let mut extractor = args.extractor();
-            let tag: &str = extractor.expect_named("tag")?;
-            extractor.check_exhausted()?;
+        fn from_args(args: &mut ArgsAccess<'_>) -> Result<Self, ExtensionError> {
+            let tag: &str = args.expect_named("tag")?;
             Ok(TestSource {
                 tag: tag.to_string(),
             })

src/extensions/args.rs

@@ -26,17 +26,17 @@
 //! into default Substrait literal expressions.
 
 use std::collections::HashSet;
+use std::fmt;
 use std::slice::Iter as SliceIter;
 use std::vec::IntoIter as VecIntoIter;
-use std::{fmt, thread};
 
 use indexmap::IndexMap;
 use substrait::proto;
 use substrait::proto::expression::field_reference::ReferenceType;
 use substrait::proto::expression::literal::LiteralType;
 use substrait::proto::expression::{RexType, reference_segment};
 
-use super::ExtensionError;
+use super::{Explainable, ExtensionError};
 use crate::textify::expressions::Reference;
 
 /// Kind of relation addendum in the text format.
@@ -198,43 +198,60 @@ pub struct ExtensionArgs {
     pub output_columns: Vec<ExtensionColumn>,
 }
 
-/// Helper struct for extracting named arguments with validation.
+/// [`ArgsAccess`] provides a view of the arguments in the text form of an
+/// extension relation or advanced extension.
 ///
-/// Tracks which arguments have been consumed. Callers **must** call
-/// [`check_exhausted`](ArgsExtractor::check_exhausted) before dropping to
-/// verify no unexpected arguments remain. In debug builds, dropping without
-/// calling `check_exhausted` will panic. This catches [`Explainable`](super::Explainable)
-/// implementations that forget to reject unexpected named arguments.
-pub struct ArgsExtractor<'a> {
+/// Positional arguments are available via [`Self::positional`], named arguments
+/// via [`Self::get_named`] or [`Self::expect_named`] which can do type
+/// conversion as well, and output columns via [`Self::output_columns`].
+///
+/// [`ArgsAccess`] tracks which arguments are accessed; when
+/// [`Explainable::from_args`] returns, any unaccessed arguments will be raised
+/// as [`ExtensionError::InvalidArgument`] errors.
+pub struct ArgsAccess<'a> {
     args: &'a ExtensionArgs,
-    consumed: HashSet<&'a str>,
-    checked: bool,
+    handled: HashSet<&'a str>,
+    positional_handled: bool,
 }
 
-impl<'a> ArgsExtractor<'a> {
-    /// Create a new extractor for the given arguments
-    pub fn new(args: &'a ExtensionArgs) -> Self {
+impl<'a> ArgsAccess<'a> {
+    pub(crate) fn new(args: &'a ExtensionArgs) -> Self {
         Self {
             args,
-            consumed: HashSet::new(),
-            checked: false,
+            handled: HashSet::new(),
+            positional_handled: false,
         }
     }
 
-    /// Get a named argument value, marking it as consumed if found.
+    /// Returns the positional arguments in source order and marks them as handled.
+    pub fn positional(&mut self) -> &'a [ExtensionValue] {
+        self.positional_handled = true;
+        &self.args.positional
+    }
+
+    /// Returns the output columns for a custom relation.
+    ///
+    /// Output columns are not included in the unhandled-argument check.
+    pub fn output_columns(&self) -> &'a [ExtensionColumn] {
+        &self.args.output_columns
+    }
+
+    /// Returns a named argument without converting it, or `None` if it is absent.
+    ///
+    /// A present argument is marked as handled.
     pub fn get_named_arg(&mut self, name: &str) -> Option<&'a ExtensionValue> {
         match self.args.named.get_key_value(name) {
             Some((k, value)) => {
-                self.consumed.insert(k);
+                self.handled.insert(k);
                 Some(value)
             }
             None => None,
         }
     }
 
-    /// Get a named argument converted to `T`, or `None` if it is absent.
+    /// Returns a named argument converted to `T`, or `None` if it is absent.
     ///
-    /// Marks the argument as consumed if it exists in the source args.
+    /// A present argument is marked as handled.
     pub fn get_named<T>(&mut self, name: &str) -> Result<Option<T>, ExtensionError>
     where
         T: TryFrom<&'a ExtensionValue>,
@@ -250,9 +267,9 @@ impl<'a> ArgsExtractor<'a> {
             .transpose()
     }
 
-    /// Get a required named argument converted to `T`.
+    /// Returns a required named argument converted to `T`.
     ///
-    /// Marks the argument as consumed if found.
+    /// A present argument is marked as handled.
     pub fn expect_named<T>(&mut self, name: &str) -> Result<T, ExtensionError>
     where
         T: TryFrom<&'a ExtensionValue>,
@@ -264,48 +281,35 @@ impl<'a> ArgsExtractor<'a> {
             })
     }
 
-    /// Check that all named arguments in the source have been consumed,
-    /// returning an error if not.
-    ///
-    /// Must be called before the extractor is dropped, to validate that all
-    /// args are correctly handled. In debug builds, dropping without calling
-    /// this method will panic.
-    pub fn check_exhausted(&mut self) -> Result<(), ExtensionError> {
-        self.checked = true;
+    /// Rejects arguments that the decoder did not handle.
+    pub(crate) fn finish(self) -> Result<(), ExtensionError> {
+        if !self.positional_handled && !self.args.positional.is_empty() {
+            return Err(ExtensionError::InvalidArgument(format!(
+                "Unhandled positional arguments: {}",
+                self.args.positional.len()
+            )));
+        }
 
-        let mut unknown_args = Vec::new();
+        let mut unhandled_args = Vec::new();
         for name in self.args.named.keys() {
-            if !self.consumed.contains(name.as_str()) {
-                unknown_args.push(name.as_str());
+            if !self.handled.contains(name.as_str()) {
+                unhandled_args.push(name.as_str());
             }
         }
 
-        if unknown_args.is_empty() {
+        if unhandled_args.is_empty() {
             Ok(())
         } else {
-            // Sort for stable error messages
-            unknown_args.sort();
+            // Sort for stable error messages.
+            unhandled_args.sort();
             Err(ExtensionError::InvalidArgument(format!(
                 "Unknown named arguments: {}",
-                unknown_args.join(", ")
+                unhandled_args.join(", ")
             )))
         }
     }
 }
 
-impl Drop for ArgsExtractor<'_> {
-    fn drop(&mut self) {
-        if self.checked || thread::panicking() {
-            return;
-        }
-        // If we get here, the caller forgot to call check_exhausted().
-        debug_assert!(
-            false,
-            "ArgsExtractor dropped without calling check_exhausted()"
-        );
-    }
-}
-
 /// A tuple-valued extension argument.
 ///
 /// Tuple values preserve positional order and can be iterated by value or by
@@ -661,6 +665,21 @@ impl ExtensionColumn {
 }
 
 impl ExtensionArgs {
+    /// Decodes these arguments as an [`Explainable`] extension value.
+    ///
+    /// If decoding succeeds, this method rejects any named or positional
+    /// arguments the implementation did not access. If decoding fails, it
+    /// returns that error without checking for unhandled arguments.
+    pub fn parse<T>(&self) -> Result<T, ExtensionError>
+    where
+        T: Explainable,
+    {
+        let mut access = ArgsAccess::new(self);
+        let value = T::from_args(&mut access)?;
+        access.finish()?;
+        Ok(value)
+    }
+
     /// Push a positional extension argument.
     pub fn push<T>(&mut self, value: T)
     where
@@ -677,11 +696,6 @@ impl ExtensionArgs {
     {
         self.named.insert(name.into(), value.into())
     }
-
-    /// Create an extractor for validating named arguments
-    pub fn extractor(&self) -> ArgsExtractor<'_> {
-        ArgsExtractor::new(self)
-    }
 }
 
 #[cfg(test)]
@@ -692,28 +706,28 @@ mod tests {
     fn get_named_converts_present_values_and_returns_none_for_missing_values() {
         let mut args = ExtensionArgs::default();
         args.insert("count", 8_i64);
-        let mut extractor = args.extractor();
+        let mut access = ArgsAccess::new(&args);
 
-        assert_eq!(extractor.get_named::<i64>("count").unwrap(), Some(8));
-        assert_eq!(extractor.get_named::<i64>("missing").unwrap(), None);
-        assert!(extractor.check_exhausted().is_ok());
+        assert_eq!(access.get_named::<i64>("count").unwrap(), Some(8));
+        assert_eq!(access.get_named::<i64>("missing").unwrap(), None);
+        assert!(access.finish().is_ok());
     }
 
     #[test]
     fn get_named_contextualizes_conversion_errors() {
         let mut args = ExtensionArgs::default();
         args.insert("count", ExtensionValue::Null);
-        let mut extractor = args.extractor();
+        let mut access = ArgsAccess::new(&args);
 
-        let error = extractor
+        let error = access
             .get_named::<i64>("count")
             .expect_err("null should not convert to i64");
 
         assert_eq!(
             error.to_string(),
             "Invalid named argument 'count': Invalid argument: expected integer, got null"
         );
-        assert!(extractor.check_exhausted().is_ok());
+        assert!(access.finish().is_ok());
     }
 
     #[test]
@@ -738,13 +752,13 @@ mod tests {
     #[test]
     fn expect_named_reports_missing_argument_name() {
         let args = ExtensionArgs::default();
-        let mut extractor = args.extractor();
+        let mut access = ArgsAccess::new(&args);
 
-        let error = extractor
+        let error = access
             .expect_named::<i64>("count")
             .expect_err("missing argument should fail");
 
         assert_eq!(error.to_string(), "Missing required argument: count");
-        assert!(extractor.check_exhausted().is_ok());
+        assert!(access.finish().is_ok());
     }
 }