Extractor rewrite

crates/pile-value/src/extract/item/flac.rs

@@ -0,0 +1,162 @@
+use mime::Mime;
+use pile_config::Label;
+use pile_flac::{FlacBlock, FlacReader};
+use std::{
+	collections::HashMap,
+	io::BufReader,
+	sync::{Arc, OnceLock},
+};
+
+use crate::{
+	extract::traits::{ListExtractor, ObjectExtractor},
+	value::{Item, PileValue, SyncReadBridge},
+};
+
+pub struct FlacImagesExtractor {
+	item: Item,
+}
+
+impl FlacImagesExtractor {
+	pub fn new(item: &Item) -> Self {
+		Self { item: item.clone() }
+	}
+
+	async fn get_images(&self) -> Result<Vec<PileValue>, std::io::Error> {
+		let reader = SyncReadBridge::new_current(self.item.read().await?);
+		let raw_images = tokio::task::spawn_blocking(move || {
+			let reader = FlacReader::new(BufReader::new(reader));
+			let mut images: Vec<(Mime, Vec<u8>)> = Vec::new();
+			for block in reader {
+				match block.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))? {
+					FlacBlock::Picture(picture) => {
+						images.push((picture.mime, picture.img_data));
+					}
+					FlacBlock::AudioFrame(_) => break,
+					_ => {}
+				}
+			}
+			Ok::<_, std::io::Error>(images)
+		})
+		.await
+		.map_err(std::io::Error::other)??;
+
+		Ok(raw_images
+			.into_iter()
+			.map(|(mime, data)| PileValue::Blob {
+				mime,
+				bytes: Arc::new(data),
+			})
+			.collect())
+	}
+}
+
+#[async_trait::async_trait]
+impl ListExtractor for FlacImagesExtractor {
+	async fn get<'a>(&'a self, idx: usize) -> Result<Option<PileValue>, std::io::Error> {
+		Ok(self.get_images().await?.into_iter().nth(idx))
+	}
+
+	async fn len(&self) -> Result<usize, std::io::Error> {
+		Ok(self.get_images().await?.len())
+	}
+}
+
+pub struct FlacExtractor {
+	item: Item,
+	output: OnceLock<HashMap<Label, PileValue>>,
+	images: Option<PileValue>,
+}
+
+impl FlacExtractor {
+	pub fn new(item: &Item) -> Self {
+		let is_flac = match item {
+			Item::File { path, .. } => path.to_str().unwrap_or_default().ends_with(".flac"),
+			Item::S3 { key, .. } => key.ends_with(".flac"),
+		};
+
+		let images =
+			is_flac.then(|| PileValue::ListExtractor(Arc::new(FlacImagesExtractor::new(item))));
+
+		Self {
+			item: item.clone(),
+			output: OnceLock::new(),
+			images,
+		}
+	}
+
+	async fn get_inner(&self) -> Result<&HashMap<Label, PileValue>, std::io::Error> {
+		if let Some(x) = self.output.get() {
+			return Ok(x);
+		}
+
+		let key = match &self.item {
+			Item::File { path, .. } => path.to_str().unwrap_or_default().to_owned(),
+			Item::S3 { key, .. } => key.to_string(),
+		};
+
+		if !key.ends_with(".flac") {
+			let _ = self.output.set(HashMap::new());
+			#[expect(clippy::unwrap_used)]
+			return Ok(self.output.get().unwrap());
+		}
+
+		let reader = SyncReadBridge::new_current(self.item.read().await?);
+		let raw_tags = tokio::task::spawn_blocking(move || {
+			let reader = FlacReader::new(BufReader::new(reader));
+			let mut tags: Vec<(String, String)> = Vec::new();
+			for block in reader {
+				match block.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))? {
+					FlacBlock::VorbisComment(comment) => {
+						for (k, v) in comment.comment.comments {
+							tags.push((k.to_string().to_lowercase(), v.into()));
+						}
+					}
+					FlacBlock::AudioFrame(_) => break,
+					_ => {}
+				}
+			}
+			Ok::<_, std::io::Error>(tags)
+		})
+		.await
+		.map_err(std::io::Error::other)??;
+
+		let mut output: HashMap<Label, Vec<PileValue>> = HashMap::new();
+		for (k, v) in raw_tags {
+			if let Some(label) = Label::new(k) {
+				output
+					.entry(label)
+					.or_default()
+					.push(PileValue::String(Arc::new(v.into())));
+			}
+		}
+		let output: HashMap<Label, PileValue> = output
+			.into_iter()
+			.map(|(k, v)| (k, PileValue::Array(Arc::new(v))))
+			.collect();
+
+		let _ = self.output.set(output);
+		#[expect(clippy::unwrap_used)]
+		return Ok(self.output.get().unwrap());
+	}
+}
+
+#[async_trait::async_trait]
+impl ObjectExtractor for FlacExtractor {
+	async fn field(&self, name: &Label) -> Result<Option<PileValue>, std::io::Error> {
+		if name.as_str() == "images"
+			&& let Some(ref images) = self.images
+		{
+			return Ok(Some(images.clone()));
+		}
+		Ok(self.get_inner().await?.get(name).cloned())
+	}
+
+	async fn fields(&self) -> Result<Vec<Label>, std::io::Error> {
+		let mut fields = self.get_inner().await?.keys().cloned().collect::<Vec<_>>();
+		if self.images.is_some() {
+			#[expect(clippy::unwrap_used)]
+			fields.push(Label::new("images").unwrap());
+		}
+		Ok(fields)
+	}
+}