Add Record::drain to take out elements by range (#11002)

Matches the general behavior of `Vec::drain` or
`indexmap::IndexMap::drain`:
- Drop the remaining elements (implementing the unstable `keep_rest()`
would not be compatible with something like `indexmap`)
- No `AsRef<[T]>` or `Drain::as_slice()` behavior as this would make
layout assumptions.
- `Drain: DoubleEndedIterator`

Found useful in #10903
This commit is contained in:
Stefan Holderbach 2023-11-08 22:54:02 +01:00 committed by GitHub
parent 1fd3bc1ba6
commit 44c0db46e1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1,3 +1,5 @@
use std::ops::RangeBounds;
use crate::Value; use crate::Value;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
@ -218,6 +220,42 @@ impl Record {
iter: self.vals.into_iter(), iter: self.vals.into_iter(),
} }
} }
/// Obtain an iterator to remove elements in `range`
///
/// Elements not consumed from the iterator will be dropped
///
/// ```rust
/// use nu_protocol::{record, Value};
///
/// let mut rec = record!(
/// "a" => Value::test_nothing(),
/// "b" => Value::test_int(42),
/// "c" => Value::test_string("foo"),
/// );
/// {
/// let mut drainer = rec.drain(1..);
/// assert_eq!(drainer.next(), Some(("b".into(), Value::test_int(42))));
/// // Dropping the `Drain`
/// }
/// let mut rec_iter = rec.into_iter();
/// assert_eq!(rec_iter.next(), Some(("a".into(), Value::test_nothing())));
/// assert_eq!(rec_iter.next(), None);
/// ```
pub fn drain<R>(&mut self, range: R) -> Drain
where
R: RangeBounds<usize> + Clone,
{
assert_eq!(
self.cols.len(),
self.vals.len(),
"Length of cols and vals must be equal for sane `Record::drain`"
);
Drain {
keys: self.cols.drain(range.clone()),
values: self.vals.drain(range),
}
}
} }
impl FromIterator<(String, Value)> for Record { impl FromIterator<(String, Value)> for Record {
@ -357,6 +395,35 @@ impl ExactSizeIterator for IntoValues {
} }
} }
pub struct Drain<'a> {
keys: std::vec::Drain<'a, String>,
values: std::vec::Drain<'a, Value>,
}
impl Iterator for Drain<'_> {
type Item = (String, Value);
fn next(&mut self) -> Option<Self::Item> {
Some((self.keys.next()?, self.values.next()?))
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.keys.size_hint()
}
}
impl DoubleEndedIterator for Drain<'_> {
fn next_back(&mut self) -> Option<Self::Item> {
Some((self.keys.next_back()?, self.values.next_back()?))
}
}
impl ExactSizeIterator for Drain<'_> {
fn len(&self) -> usize {
self.keys.len()
}
}
#[macro_export] #[macro_export]
macro_rules! record { macro_rules! record {
{$($col:expr => $val:expr),+ $(,)?} => { {$($col:expr => $val:expr),+ $(,)?} => {