spire_core/dataset/
data.rs

1use std::fmt;
2use std::future::Future;
3use std::sync::Arc;
4
5use futures::future::BoxFuture;
6use futures::{FutureExt, Stream, stream};
7
8use crate::Result;
9use crate::dataset::Dataset;
10
11/// A cloneable, type-erased handle to a [`Dataset`].
12pub struct Data<T> {
13    inner: Arc<dyn DynDataset<T>>,
14}
15
16impl<T: Send + 'static> Data<T> {
17    /// Wraps `dataset`.
18    #[must_use]
19    pub fn new(dataset: impl Dataset<T>) -> Self {
20        Self {
21            inner: Arc::new(dataset),
22        }
23    }
24
25    /// Stores an item.
26    ///
27    /// # Errors
28    ///
29    /// Returns an error if the underlying dataset fails.
30    pub async fn write(&self, item: T) -> Result<()> {
31        self.inner.write(item).await
32    }
33
34    /// Stores every item of `items`, stopping at the first error.
35    ///
36    /// # Errors
37    ///
38    /// Returns an error if the underlying dataset fails.
39    pub async fn write_all(&self, items: impl IntoIterator<Item = T>) -> Result<()> {
40        for item in items {
41            self.inner.write(item).await?;
42        }
43        Ok(())
44    }
45
46    /// Takes the next item, or `None` if the dataset is empty.
47    ///
48    /// # Errors
49    ///
50    /// Returns an error if the underlying dataset fails.
51    pub async fn read(&self) -> Result<Option<T>> {
52        self.inner.read().await
53    }
54
55    /// Takes every item until the dataset is empty.
56    ///
57    /// # Errors
58    ///
59    /// Returns an error if the underlying dataset fails.
60    pub async fn read_all(&self) -> Result<Vec<T>> {
61        let mut items = Vec::with_capacity(self.size_hint().unwrap_or_default());
62        while let Some(item) = self.inner.read().await? {
63            items.push(item);
64        }
65        Ok(items)
66    }
67
68    /// Returns a stream that takes items until the dataset is empty.
69    pub fn into_stream(self) -> impl Stream<Item = Result<T>> + Send {
70        stream::try_unfold(self, |data| async move {
71            let item = data.read().await?;
72            Ok(item.map(|item| (item, data)))
73        })
74    }
75
76    /// Returns the number of stored items, if known.
77    #[must_use]
78    pub fn size_hint(&self) -> Option<usize> {
79        self.inner.size_hint()
80    }
81}
82
83impl<T> Clone for Data<T> {
84    fn clone(&self) -> Self {
85        Self {
86            inner: self.inner.clone(),
87        }
88    }
89}
90
91impl<T> fmt::Debug for Data<T> {
92    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
93        f.debug_struct("Data")
94            .field("type", &std::any::type_name::<T>())
95            .finish_non_exhaustive()
96    }
97}
98
99impl<T: Send + 'static> Dataset<T> for Data<T> {
100    fn write(&self, item: T) -> impl Future<Output = Result<()>> + Send {
101        self.inner.write(item)
102    }
103
104    fn read(&self) -> impl Future<Output = Result<Option<T>>> + Send {
105        self.inner.read()
106    }
107
108    fn size_hint(&self) -> Option<usize> {
109        self.inner.size_hint()
110    }
111}
112
113/// An object-safe [`Dataset`].
114trait DynDataset<T>: Send + Sync + 'static {
115    fn write(&self, item: T) -> BoxFuture<'_, Result<()>>;
116    fn read(&self) -> BoxFuture<'_, Result<Option<T>>>;
117    fn size_hint(&self) -> Option<usize>;
118}
119
120impl<T, D> DynDataset<T> for D
121where
122    T: Send + 'static,
123    D: Dataset<T>,
124{
125    fn write(&self, item: T) -> BoxFuture<'_, Result<()>> {
126        Dataset::write(self, item).boxed()
127    }
128
129    fn read(&self) -> BoxFuture<'_, Result<Option<T>>> {
130        Dataset::read(self).boxed()
131    }
132
133    fn size_hint(&self) -> Option<usize> {
134        Dataset::size_hint(self)
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use futures::TryStreamExt;
141
142    use super::*;
143    use crate::dataset::InMemDataset;
144
145    #[tokio::test]
146    async fn batch_and_stream() {
147        let data = Data::new(InMemDataset::new());
148        data.write_all([1, 2, 3]).await.unwrap();
149        assert_eq!(data.size_hint(), Some(3));
150        assert_eq!(data.read().await.unwrap(), Some(1));
151        assert_eq!(
152            data.clone()
153                .into_stream()
154                .try_collect::<Vec<_>>()
155                .await
156                .unwrap(),
157            [2, 3]
158        );
159        assert_eq!(data.read_all().await.unwrap(), Vec::<i32>::new());
160    }
161}
spire_core/dataset/data.rs

spire_core/dataset/
data.rs
