Add `at` method to `arguments` (#4864)

2 files changed, 91 insertions, 2 deletions

diff --git a/crates/typst/src/foundations/args.rs b/crates/typst/src/foundations/args.rs
index d580be3c..db443648 100644
--- a/crates/typst/src/foundations/args.rs
+++ b/crates/typst/src/foundations/args.rs
@@ -2,9 +2,9 @@ use std::fmt::{self, Debug, Formatter};
 
 use ecow::{eco_format, eco_vec, EcoString, EcoVec};
 
-use crate::diag::{bail, error, At, SourceDiagnostic, SourceResult};
+use crate::diag::{bail, error, At, SourceDiagnostic, SourceResult, StrResult};
 use crate::foundations::{
-    func, repr, scope, ty, Array, Dict, FromValue, IntoValue, Repr, Str, Value,
+    cast, func, repr, scope, ty, Array, Dict, FromValue, IntoValue, Repr, Str, Value,
 };
 use crate::syntax::{Span, Spanned};
 
@@ -256,6 +256,42 @@ impl Args {
     }
 }
 
+/// A key that can be used to get an argument: either the index of a positional
+/// argument, or the name of a named argument.
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum ArgumentKey {
+    Index(i64),
+    Name(Str),
+}
+
+cast! {
+    ArgumentKey,
+    v: i64 => Self::Index(v),
+    v: Str => Self::Name(v),
+}
+
+impl Args {
+    fn get(&self, key: &ArgumentKey) -> Option<&Value> {
+        let item = match key {
+            &ArgumentKey::Index(index) => {
+                let mut iter = self.items.iter().filter(|item| item.name.is_none());
+                if index < 0 {
+                    let index = (-(index + 1)).try_into().ok()?;
+                    iter.nth_back(index)
+                } else {
+                    let index = index.try_into().ok()?;
+                    iter.nth(index)
+                }
+            }
+            // Accept the last argument with the right name.
+            ArgumentKey::Name(name) => {
+                self.items.iter().rfind(|item| item.name.as_ref() == Some(name))
+            }
+        };
+        item.map(|item| &item.value.v)
+    }
+}
+
 #[scope]
 impl Args {
     /// Construct spreadable arguments in place.
@@ -279,6 +315,28 @@ impl Args {
         args.take()
     }
 
+    /// Returns the positional argument at the specified index, or the named
+    /// argument with the specified name.
+    ///
+    /// If the key is an [integer]($int), this is equivalent to first calling
+    /// [`pos`]($arguments.pos) and then [`array.at`]. If it is a [string]($str),
+    /// this is equivalent to first calling [`named`]($arguments.named) and then
+    /// [`dictionary.at`].
+    #[func]
+    pub fn at(
+        &self,
+        /// The index or name of the argument to get.
+        key: ArgumentKey,
+        /// A default value to return if the key is invalid.
+        #[named]
+        default: Option<Value>,
+    ) -> StrResult<Value> {
+        self.get(&key)
+            .cloned()
+            .or(default)
+            .ok_or_else(|| missing_key_no_default(key))
+    }
+
     /// Returns the captured positional arguments as an array.
     #[func(name = "pos", title = "Positional")]
     pub fn to_pos(&self) -> Array {
@@ -380,3 +438,16 @@ where
         Args::new(fallback, self)
     }
 }
+
+/// The missing key access error message when no default was given.
+#[cold]
+fn missing_key_no_default(key: ArgumentKey) -> EcoString {
+    eco_format!(
+        "arguments do not contain key {} \
+         and no default value was specified",
+        match key {
+            ArgumentKey::Index(i) => i.repr(),
+            ArgumentKey::Name(name) => name.repr(),
+        }
+    )
+}
diff --git a/tests/suite/scripting/arguments.typ b/tests/suite/scripting/arguments.typ
new file mode 100644
index 00000000..e82f4962
--- /dev/null
+++ b/tests/suite/scripting/arguments.typ
@@ -0,0 +1,18 @@
+// Test arguments.
+
+--- arguments-at ---
+#let args = arguments(0, 1, a: 2, 3)
+#test(args.at(0), 0)
+#test(args.at(1), 1)
+#test(args.at(2), 3)
+#test(args.at("a"), 2)
+
+--- arguments-at-invalid-index ---
+#let args = arguments(0, 1, a: 2, 3)
+// Error: 2-12 arguments do not contain key 4 and no default value was specified
+#args.at(4)
+
+--- arguments-at-invalid-name ---
+#let args = arguments(0, 1, a: 2, 3)
+// Error: 2-14 arguments do not contain key "b" and no default value was specified
+#args.at("b")