nolgia 0.1.0

Official Rust SDK for the Nolgia API: image, video, audio and agent generation, generated from the OpenAPI spec
Documentation
//! Hand-written. Cursor pagination as a [`Stream`] of items.

use std::{collections::VecDeque, future::Future};

use futures_util::{Stream, stream};

use crate::{
    Client, Typed,
    api::{builder, types},
    error::Result,
};

/// One page from a cursor-paginated list.
#[derive(Debug, Clone)]
pub struct Page<T> {
    /// The page's items.
    pub items: Vec<T>,
    /// The cursor for the next page; `None` on the last one.
    pub next_cursor: Option<String>,
}

impl From<types::AssetPage> for Page<types::Asset> {
    fn from(page: types::AssetPage) -> Self {
        Self {
            items: page.items,
            next_cursor: page.next_cursor,
        }
    }
}

impl From<types::JobPage> for Page<types::Job> {
    fn from(page: types::JobPage) -> Self {
        Self {
            items: page.items,
            next_cursor: page.next_cursor,
        }
    }
}

/// Every item of a cursor-paginated list, fetching pages as the stream is read.
///
/// `fetch` gets `None` for the first page and each `next_cursor` after it.
/// Use it for any list operation; [`Client::paginate_assets`] and
/// [`Client::paginate_jobs`] are ready-made.
pub fn paginate<T, F, Fut>(fetch: F) -> impl Stream<Item = Result<T>>
where
    F: FnMut(Option<String>) -> Fut,
    Fut: Future<Output = Result<Page<T>>>,
{
    struct State<T, F> {
        fetch: F,
        cursor: Option<String>,
        buffer: VecDeque<T>,
        done: bool,
    }
    stream::unfold(
        State {
            fetch,
            cursor: None,
            buffer: VecDeque::new(),
            done: false,
        },
        |mut state| async move {
            loop {
                if let Some(item) = state.buffer.pop_front() {
                    return Some((Ok(item), state));
                }
                if state.done {
                    return None;
                }
                match (state.fetch)(state.cursor.take()).await {
                    Ok(page) => {
                        state.buffer.extend(page.items);
                        state.cursor = page.next_cursor.filter(|c| !c.is_empty());
                        state.done = state.cursor.is_none();
                    }
                    Err(err) => {
                        state.done = true;
                        return Some((Err(err), state));
                    }
                }
            }
        },
    )
}

impl Client {
    /// Every asset matching the filters you set on the list builder.
    ///
    /// ```no_run
    /// # async fn run(nolgia: nolgia::Client) -> nolgia::Result<()> {
    /// use nolgia::{StreamExt, types::Modality};
    ///
    /// let videos = nolgia.paginate_assets(|list| list.modality(Modality::Video).limit(100u64));
    /// let mut videos = std::pin::pin!(videos);
    /// while let Some(asset) = videos.next().await {
    ///     println!("{}", asset?.id);
    /// }
    /// # Ok(()) }
    /// ```
    pub fn paginate_assets<'a>(
        &'a self,
        configure: impl Fn(builder::ListAssets<'a>) -> builder::ListAssets<'a> + 'a,
    ) -> impl Stream<Item = Result<types::Asset>> + 'a {
        paginate(move |cursor| {
            let mut list = configure(self.api.list_assets());
            if let Some(cursor) = cursor {
                list = list.cursor(cursor);
            }
            async move { list.send().typed().await.map(Page::from) }
        })
    }

    /// Every job matching the filters you set on the list builder.
    pub fn paginate_jobs<'a>(
        &'a self,
        configure: impl Fn(builder::ListJobs<'a>) -> builder::ListJobs<'a> + 'a,
    ) -> impl Stream<Item = Result<types::Job>> + 'a {
        paginate(move |cursor| {
            let mut list = configure(self.api.list_jobs());
            if let Some(cursor) = cursor {
                list = list.cursor(cursor);
            }
            async move { list.send().typed().await.map(Page::from) }
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures_util::StreamExt;

    #[tokio::test]
    async fn follows_cursors_to_the_end() {
        let pages = paginate(|cursor: Option<String>| async move {
            Ok(match cursor.as_deref() {
                None => Page {
                    items: vec![1, 2],
                    next_cursor: Some("b".into()),
                },
                Some("b") => Page {
                    items: vec![3],
                    next_cursor: Some(String::new()),
                },
                Some(other) => panic!("unexpected cursor {other}"),
            })
        });
        let items: Vec<i32> = pages.map(|r| r.unwrap()).collect().await;
        assert_eq!(items, vec![1, 2, 3]);
    }
}