Instantly share code, notes, and snippets.

Embed
What would you like to do?
Python list comprehensions in Rust
#![feature(macro_rules)]
/// MIT license etc. etc.
///
/// Experimenting with Python-like list comprehensions.
///
/// An attempt to explore the possibility space for ergonomic improvements
/// in Rust that can come post v1.0. Notice that there are not type declerations.
/// Aside from the "c!" macro invocation, Rust allows for an exact copy of the
/// Python comprehension syntax.
///
/// This should probably work on slices. Something to look forward to.
use std::num::{abs, pow}; // Used in unit tests only.
macro_rules! c(
($m:expr for $el:ident in $iter:expr) => (
{
let (lower, upper) = $iter.size_hint();
let mut result = match upper {
Some(bound) => Vec::with_capacity(bound),
None => Vec::with_capacity(lower),
};
for $el in $iter {
let $el = $el.clone();
result.push($m);
}
result
}
);
($m:expr for $el:ident in $iter:expr if $cond:expr) => (
{
let mut result = Vec::new();
for $el in $iter {
let $el = $el.clone();
if $cond {
result.push($m);
}
}
result
}
);
($m:expr for $el:ident in $iter:expr
$(for $inner_el:ident in $inner_iter:expr)*
$(if $cond:expr)*) => (
{
let mut result = Vec::new();
for $el in $iter {
let $el = $el.clone();
result.push_all_move(c!($($m for $inner_el in $inner_iter)+ $(if $cond)*));
}
result
}
);
)
#[test]
fn simple_comprehension() {
let squares = c![ x * x for x in range(0, 10i) ];
assert_eq!(squares, vec![0i, 1, 4, 9, 16, 25, 36, 49, 64, 81]);
}
#[test]
fn vector_to_iter_comprehension() {
let vec = vec![-4i, -2, 0, 2, 4];
let output: Vec<int> = c![x*2 for x in vec.iter()];
assert_eq!(output, vec![-8i, -4, 0, 4, 8]);
}
#[test]
fn filter_comprehension() {
let vec = vec![-4i, -2, 0, 2, 4];
let output = c![x for x in vec.iter() if x >= 0i];
assert_eq!(output, vec![0i, 2, 4]);
}
#[test]
fn apply_function_comprehension() {
let vec = vec![-4i, -2, 0, 2, 4];
let output = c![abs(x) for x in vec.iter()];
assert_eq!(output, vec![4i, 2, 0, 2, 4]);
}
#[test]
fn call_method_comprehension() {
let fresh_fruit = vec![" banana", " loganberry ", "passion fruit "];
let output = c![fruit.trim_chars(' ') for fruit in fresh_fruit.iter()];
assert_eq!(output, vec!["banana", "loganberry", "passion fruit"]);
}
#[test]
fn change_type_comprehension() {
let output = c![(x, pow(x,2)) for x in range(0i,6)];
assert_eq!(output, vec![(0i, 0i), (1, 1), (2, 4), (3, 9), (4, 16), (5, 25)]);
}
#[test]
fn combined_comprehension() {
let a = vec![1i,2,3];
let b = vec![3i,1,4];
let not_equal = c![(x, y) for x in a.iter() for y in b.iter() if x != y];
assert_eq!(not_equal, vec![(1i, 3i), (1, 4), (2, 3), (2, 1), (2, 4), (3, 1), (3, 4)]);
println!("{}", a);
}
#[test]
fn nested_comprehensions() {
let matrix = vec![vec![1i,2,3,4],vec![5,6,7,8],vec![9,10,11,12]];
let output = c![ c![ row[i] for row in matrix.iter()] for i in range(0u,4)];
assert_eq!(output, vec![vec![1i,5,9], vec![2i,6,10], vec![3i,7,11], vec![4i,8,12]]);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment