Add exact argument to array.zip

crates/typst/src/foundations/array.rs

@@ -8,14 +8,14 @@ use ecow::{eco_format, EcoString, EcoVec};
 use serde::{Deserialize, Serialize};
 use smallvec::SmallVec;
 
-use crate::diag::{bail, At, SourceResult, StrResult};
+use crate::diag::{bail, At, SourceDiagnostic, SourceResult, StrResult};
 use crate::engine::Engine;
 use crate::eval::ops;
 use crate::foundations::{
  cast, func, repr, scope, ty, Args, Bytes, CastInfo, Context, Dict, FromValue, Func,
  IntoValue, Reflect, Repr, Str, Value, Version,
 };
-use crate::syntax::Span;
+use crate::syntax::{Span, Spanned};
 
 /// Create a new [`Array`] from values.
 #[macro_export]
@@ -482,9 +482,14 @@ impl Array {
  #[func]
  pub fn zip(
  self,
- /// The real arguments (the other arguments are just for the docs, this
- /// function is a bit involved, so we parse the arguments manually).
+ /// The real arguments (the `others` arguments are just for the docs, this
+ /// function is a bit involved, so we parse the positional arguments manually).
  args: &mut Args,
+ /// Whether all arrays have to have the same length.
+ /// For example, `(1, 2).zip((1, 2, 3), exact: true)` produces an error.
+ #[named]
+ #[default(false)]
+ exact: bool,
  /// The arrays to zip with.
  #[external]
  #[variadic]
@@ -499,7 +504,16 @@ impl Array {
 
  // Fast path for just two arrays.
  if remaining == 1 {
- let other = args.expect::<Array>("others")?;
+ let Spanned { v: other, span: other_span } =
+ args.expect::<Spanned<Array>>("others")?;
+ if exact && self.len() != other.len() {
+ bail!(
+ other_span,
+ "second array has different length ({}) from first array ({})",
+ other.len(),
+ self.len()
+ );
+ }
  return Ok(self
  .into_iter()
  .zip(other)
@@ -509,11 +523,29 @@ impl Array {
 
  // If there is more than one array, we use the manual method.
  let mut out = Self::with_capacity(self.len());
- let mut iterators = args
- .all::<Array>()?
- .into_iter()
- .map(|i| i.into_iter())
- .collect::<Vec<_>>();
+ let arrays = args.all::<Spanned<Array>>()?;
+ if exact {
+ let errs = arrays
+ .iter()
+ .filter(|sp| sp.v.len() != self.len())
+ .map(|Spanned { v, span }| {
+ SourceDiagnostic::error(
+ *span,
+ eco_format!(
+ "array has different length ({}) from first array ({})",
+ v.len(),
+ self.len()
+ ),
+ )
+ })
+ .collect::<EcoVec<_>>();
+ if !errs.is_empty() {
+ return Err(errs);
+ }
+ }
+
+ let mut iterators =
+ arrays.into_iter().map(|i| i.v.into_iter()).collect::<Vec<_>>();
 
  for this in self {
  let mut row = Self::with_capacity(1 + iterators.len());

tests/suite/foundations/array.typ

@@ -350,6 +350,7 @@
 #test((1,).zip(()), ())
 #test((1,).zip((2,)), ((1, 2),))
 #test((1, 2).zip((3, 4)), ((1, 3), (2, 4)))
+#test((1, 2).zip((3, 4), exact: true), ((1, 3), (2, 4)))
 #test((1, 2, 3, 4).zip((5, 6)), ((1, 5), (2, 6)))
 #test(((1, 2), 3).zip((4, 5)), (((1, 2), 4), (3, 5)))
 #test((1, "hi").zip((true, false)), ((1, true), ("hi", false)))
@@ -359,6 +360,15 @@
 #test((1, 2, 3).zip(), ((1,), (2,), (3,)))
 #test(array.zip(()), ())
 
+--- array-zip-exact-error ---
+// Error: 13-22 second array has different length (3) from first array (2)
+#(1, 2).zip((1, 2, 3), exact: true)
+
+--- array-zip-exact-multi-error ---
+// Error: 13-22 array has different length (3) from first array (2)
+// Error: 24-36 array has different length (4) from first array (2)
+#(1, 2).zip((1, 2, 3), (1, 2, 3, 4), exact: true)
+
 --- array-enumerate ---
 // Test the `enumerate` method.
 #test(().enumerate(), ())