74 lines
1.8 KiB
Rust
74 lines
1.8 KiB
Rust
use itertools::Itertools;
|
|
use std::ops::{Range, RangeFrom, RangeFull, RangeTo};
|
|
|
|
// This trait is already contained in the standard lib, however it is unstable.
|
|
// TODO: Remove when the `collections_range` feature stabilises
|
|
// (https://github.com/rust-lang/rust/issues/30877)
|
|
pub trait RangeArgument<T: ?Sized> {
|
|
fn start(&self) -> Option<&T>;
|
|
fn end(&self) -> Option<&T>;
|
|
}
|
|
|
|
impl<T: ?Sized> RangeArgument<T> for RangeFull {
|
|
fn start(&self) -> Option<&T> {
|
|
None
|
|
}
|
|
fn end(&self) -> Option<&T> {
|
|
None
|
|
}
|
|
}
|
|
|
|
impl<T> RangeArgument<T> for RangeFrom<T> {
|
|
fn start(&self) -> Option<&T> {
|
|
Some(&self.start)
|
|
}
|
|
fn end(&self) -> Option<&T> {
|
|
None
|
|
}
|
|
}
|
|
|
|
impl<T> RangeArgument<T> for RangeTo<T> {
|
|
fn start(&self) -> Option<&T> {
|
|
None
|
|
}
|
|
fn end(&self) -> Option<&T> {
|
|
Some(&self.end)
|
|
}
|
|
}
|
|
|
|
impl<T> RangeArgument<T> for Range<T> {
|
|
fn start(&self) -> Option<&T> {
|
|
Some(&self.start)
|
|
}
|
|
fn end(&self) -> Option<&T> {
|
|
Some(&self.end)
|
|
}
|
|
}
|
|
|
|
/// Take a range of lines from a string.
|
|
pub fn take_lines<R: RangeArgument<usize>>(s: &str, range: R) -> String {
|
|
let start = *range.start().unwrap_or(&0);
|
|
let mut lines = s.lines().skip(start);
|
|
match range.end() {
|
|
Some(&end) => lines.take(end.saturating_sub(start)).join("\n"),
|
|
None => lines.join("\n"),
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::take_lines;
|
|
|
|
#[test]
|
|
fn take_lines_test() {
|
|
let s = "Lorem\nipsum\ndolor\nsit\namet";
|
|
assert_eq!(take_lines(s, 1..3), "ipsum\ndolor");
|
|
assert_eq!(take_lines(s, 3..), "sit\namet");
|
|
assert_eq!(take_lines(s, ..3), "Lorem\nipsum\ndolor");
|
|
assert_eq!(take_lines(s, ..), s);
|
|
// corner cases
|
|
assert_eq!(take_lines(s, 4..3), "");
|
|
assert_eq!(take_lines(s, ..100), s);
|
|
}
|
|
}
|