How can I extend a trait in Rust defined somewhere else?

Viewed 62

I want to create a new iterator method such as:

let test_data = vec![1,2,3,1,1,1,1];
let indexes_with_val_1 = test_data.iter().find_all(|element| element == 1).unwrap();

assert_eq!(indexes_with_val_1, vec!(0,3,4,5,6));

So essentially I want to add a new method to the std::iter::Iterator trait but can't find examples that work for this

The logic is not an issue as I have a free function which works fine, I just would like to be able to use it as I have in the code example for better ergonomics.

2 Answers

You can use a design pattern called extension traits. You can't extend the Iterator trait, but you can write a new one. Here's what we're going to do.

  1. Write a new trait, IteratorExt, which has your custom method in it.
  2. Write a blanket impl that implements IteratorExt for any type that implements Iterator.
  3. Import IteratorExt to get access to your extension function.

For example, we can add a simple function called my_extension to iterators like so

trait IteratorExt {

  fn my_extension(self) -> Self;

}

impl<T: Iterator> IteratorExt for T {

  fn my_extension(self) -> Self {
    println!("Hey, it worked!");
    self
  }

}

pub fn main() {
  let x = vec!(1, 2, 3, 4);
  let y = x.iter().my_extension().map(|x| x + 1).collect::<Vec<_>>();
  println!("{:?}", y);
}

The only downside is that you have to import the new trait to use it. So if you want to use my_extension in another file, you have to import IteratorExt specifically in order to do so.

In my experience, the Rust community is somewhat divided on whether this is legitimate practice or whether it's a hack to be avoided, so your mileage may vary.

There exist the extension trait pattern for that.

The idea is that you create a trait, usually named TraitExt by convention, and implement it for all implementations of Trait:

pub trait IteratorExt {
    fn my_iterator_extension(&self);
}

impl<I: Iterator + ?Sized> IteratorExt for I {
    fn my_iterator_extension(&self) {
        // Do work.
    }
}

my_iterator.my_iterator_extension();
Related