imdb-async 0.1.0

Opinionated and unopinionated async wrappers to efficiently retrieve and parse IMDB's dataset
Documentation
use std::collections::HashMap;
use std::collections::HashSet;

use csv_async::AsyncReaderBuilder;

use futures::future::ready;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;

use serde::de::Unexpected;
use serde::Deserialize;
use serde::Deserializer;

use crate::Error;
use crate::Episode;
use crate::Movie;
use crate::Show;
use crate::TITLES_URL;
use crate::get_episodes_filtered;
use crate::get_movies_filtered;
use crate::get_shows_filtered;
use crate::start_stream;
use crate::start_stream_lines;

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Title {
	#[serde(rename = "tconst", deserialize_with = "crate::util::parse_imdb_id")]
	pub imdb_id: u64,
	pub title_type: TitleType,
	//_title_type_string: String, // tsv seems to have a bug where it doesn't remove the TitleType value after we deserialize it.  It works when we use JSON.
	pub primary_title: String,
	pub original_title: String,
	pub is_adult: u8,
	#[serde(deserialize_with = "crate::util::parse_janky_tsv_option")]
	pub start_year: Option<u16>,
	#[serde(deserialize_with = "crate::util::parse_janky_tsv_option")]
	pub end_year: Option<u16>,
	#[serde(deserialize_with = "crate::util::parse_janky_tsv_option")]
	pub runtime_minutes: Option<u16>,
	#[serde(with = "serde_with::StringWithSeparator::<serde_with::CommaSeparator>")]
	pub genres: Vec<String>
}

#[derive(Debug, Eq, PartialEq)]
pub enum TitleType {
	Short,
	Movie,
	Episode,
	TVMiniSeries,
	TVSeries,
	TVShort,
	TVSpecial,
	Video,
	VideoGame,
	Audiobook,
	RadioSeries
}

impl<'de> Deserialize<'de> for TitleType {
	fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
		let s: &str = Deserialize::deserialize(deserializer)?;
		Ok(match s {
			"short" => Self::Short,
			"episode" | "tvEpisode" => Self::Episode,
			"tvMiniSeries" => Self::TVMiniSeries,
			"movie" | "tvMovie" => Self::Movie,
			"tvSeries" => Self::TVSeries,
			"tvShort" => Self::TVShort,
			"tvSpecial" => Self::TVSpecial,
			"video" => Self::Video,
			"videoGame" => Self::VideoGame,
			"audiobook" => Self::Audiobook,
			"radioSeries" => Self::RadioSeries,
			s => return Err(serde::de::Error::invalid_value(Unexpected::Str(&s), &"'short', 'movie', 'tvEpisode', 'tvMiniSeries', 'tvMovie', 'tvSeries', 'tvShort', 'tvSpecial', 'video', 'videoGame'"))
		})
	}
}

// TODO:  Rejigger this whole shebang so that we don't need to duplicate most
//    of the code from both get_movies_filtered() and get_shows_filtered(); in
//    all likelihood, this means we pass a closure down to the individual
//    stream_*_filtered() functions.
pub async fn get_titles_filtered_split(movie_ids: &[u64], show_ids: &[u64]) -> Result<(Vec<Movie>, Vec<Show>), Error> {
	if(show_ids.len() == 0) {
		return Ok((get_movies_filtered(movie_ids).await?, vec![]));
	} else if(movie_ids.len() == 0) {
		return Ok((vec![], get_shows_filtered(show_ids).await?));
	}

	let mut episode_links = get_episodes_filtered(show_ids).await?;
	let mut all_ids = Vec::from(movie_ids);
	all_ids.extend(show_ids);
	all_ids.extend(episode_links.iter().map(|(id, _)| id));

	let mut stream = stream_titles_filtered(&all_ids).await?;
	let mut episodes_by_show = HashMap::new();
	let mut movies = Vec::with_capacity(movie_ids.len());
	let mut shows = Vec::with_capacity(show_ids.len());
	while let Some(result) = stream.next().await {
		let title = result?;
		if let Some(link) = episode_links.remove(&title.imdb_id) {
			let series_id = link.series_imdb_id;
			if let Some(result) = Episode::from_title_and_link(title, link) {
				let episode = result?;
				let episodes = episodes_by_show.entry(series_id).or_insert_with(|| Vec::new());
				episodes.push(episode);
			}
		} else if let Some(show) = Show::from_title(&title) {
			shows.push(show?);
		} else if let Some(movie) = Movie::from_title(&title) {
			movies.push(movie?);
		}
	}

	for show in shows.iter_mut() {
		show.episodes = match episodes_by_show.remove(&show.imdb_id) {
			Some(v) => v,
			None => Vec::new()
		};
	}

	Ok((movies, shows))
}

pub async fn get_titles_filtered(ids: &[u64]) -> Result<Vec<Title>, Error> {
	let ids = Vec::from(ids);
	let mut stream = stream_titles_filtered(&ids).await?;
	let mut titles = Vec::with_capacity(ids.len());
	while let Some(result) = stream.next().await {
		titles.push(result?);
	}
	Ok(titles)
}

pub async fn stream_titles_filtered(ids: &[u64]) -> Result<impl Stream<Item = Result<Title, Error>> + '_, Error> {
	let mut ids: HashSet<String> = ids.iter().map(|id| format!("tt{:07}", id)).collect();
	let reader = start_stream_lines(TITLES_URL).await?;
	let stream = AsyncReaderBuilder::new()
		.delimiter(b'\t')
		.has_headers(false)
		.create_deserializer(reader
			.filter(move |result| match result {
				Ok(line) => {
					// This is split out into multiple lines currently to make it easier to add everybody's favorite - printf debugging!
					if(ids.remove(line.split('\t').nth(0).unwrap())) {
						ready(true)
					} else {
						ready(false)
					}
				}
				Err(_) => ready(true)
			})
			.map_ok(|l| l + "\n")
			.into_async_read()
		)
		.into_deserialize::<Title>();
	Ok(stream
		//.map_err(|e| Error::from(e))
		/*
		.filter(move |result| match result {
			//Ok(line) => ready(ids.contains(line.split('\t').nth(0).unwrap())),
			Ok(title) => ready(ids.contains(&title.imdb_id)),
			Err(_) => ready(true)
		})
		*/
		//.map(|r| r.flatten())
		.err_into()
	)
}

pub async fn stream_titles() -> Result<impl Stream<Item = Result<Title, Error>>, Error> {
	let reader = start_stream(TITLES_URL).await?;
	let stream = AsyncReaderBuilder::new()
		.delimiter(b'\t')
		.create_deserializer(reader)
		.into_deserialize::<Title>();
	Ok(stream
		//.map_err(|e| Error::from(e))
		//.map_ok(move |s| tsv::de::from_str(&s).map_err(|e| Error::from(e)))
		//.map(|r| r.flatten())
		.err_into()
	)
}