Skip to main content

lumen_dataset/transform/
map.rs

1use crate::Dataset;
2
3pub trait Map {
4    type Item;
5    type Output;
6    fn map(&self, item: Self::Item) -> Self::Output;
7}
8
9pub struct MapDataset<D, M> 
10where 
11    D: Dataset,
12    M: Map<Item = D::Item>,
13{
14    dataset: D,
15    map: M,
16}
17
18impl<D, M> MapDataset<D, M> 
19where 
20    D: Dataset,
21    M: Map<Item = D::Item>,
22{
23    pub fn new(dataset: D, map: M,) -> Self {
24        Self { dataset, map }
25    }
26}
27
28impl<D, M> Dataset for MapDataset<D, M> 
29where 
30    D: Dataset,
31    M: Map<Item = D::Item>,
32{
33    type Item = M::Output;
34    type Error = D::Error;
35
36    fn get(&self, index: usize) -> Result<Option<Self::Item>, Self::Error> {
37        let item = self.dataset.get(index)?;
38        Ok(item.map(|item| self.map.map(item)))
39    }
40
41    fn len(&self) -> usize {
42        self.dataset.len()   
43    }
44}
45