| 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); |
| } |
| } |