use crate::common::PIPESIZE; use crate::meta_plugin::{ process_metadata_outputs, register_meta_plugin, BaseMetaPlugin, MetaPlugin, MetaPluginResponse, MetaPluginType, }; #[derive(Debug, Default)] pub struct InferMetaPlugin { buffer: Vec, max_buffer_size: usize, is_finalized: bool, base: BaseMetaPlugin, } impl InferMetaPlugin { pub fn new( options: Option>, outputs: Option>, ) -> InferMetaPlugin { let mut base = BaseMetaPlugin::new(); if let Some(opts) = options { for (key, value) in opts { base.options.insert(key, value); } } let max_buffer_size = base .options .get("max_buffer_size") .and_then(|v| v.as_u64()) .unwrap_or(PIPESIZE as u64) as usize; base.outputs.insert( "infer_mime_type".to_string(), serde_yaml::Value::String("infer_mime_type".to_string()), ); if let Some(outs) = outputs { for (key, value) in outs { base.outputs.insert(key, value); } } InferMetaPlugin { buffer: Vec::new(), max_buffer_size, is_finalized: false, base, } } } impl MetaPlugin for InferMetaPlugin { fn meta_type(&self) -> MetaPluginType { MetaPluginType::Infer } fn is_finalized(&self) -> bool { self.is_finalized } fn set_finalized(&mut self, finalized: bool) { self.is_finalized = finalized; } fn set_save_meta(&mut self, save_meta: crate::meta_plugin::SaveMetaFn) { self.base.set_save_meta(save_meta); } fn save_meta(&self, name: &str, value: &str) { self.base.save_meta(name, value); } fn update(&mut self, data: &[u8]) -> MetaPluginResponse { if self.is_finalized { return MetaPluginResponse { metadata: Vec::new(), is_finalized: true, }; } let remaining = self.max_buffer_size.saturating_sub(self.buffer.len()); let to_add = &data[..data.len().min(remaining)]; self.buffer.extend_from_slice(to_add); if self.buffer.len() >= self.max_buffer_size { let mime_type = infer::get(&self.buffer) .map(|kind| kind.mime_type().to_string()) .unwrap_or_else(|| "application/octet-stream".to_string()); self.is_finalized = true; let metadata = process_metadata_outputs( "infer_mime_type", serde_yaml::Value::String(mime_type), self.base.outputs(), ) .map(|m| vec![m]) .unwrap_or_default(); return MetaPluginResponse { metadata, is_finalized: true, }; } MetaPluginResponse { metadata: Vec::new(), is_finalized: false, } } fn finalize(&mut self) -> MetaPluginResponse { if self.is_finalized { return MetaPluginResponse { metadata: Vec::new(), is_finalized: true, }; } let mime_type = infer::get(&self.buffer) .map(|kind| kind.mime_type().to_string()) .unwrap_or_else(|| "application/octet-stream".to_string()); self.is_finalized = true; let metadata = process_metadata_outputs( "infer_mime_type", serde_yaml::Value::String(mime_type), self.base.outputs(), ) .map(|m| vec![m]) .unwrap_or_default(); MetaPluginResponse { metadata, is_finalized: true, } } fn outputs(&self) -> &std::collections::HashMap { self.base.outputs() } fn outputs_mut( &mut self, ) -> anyhow::Result<&mut std::collections::HashMap> { Ok(self.base.outputs_mut()) } fn default_outputs(&self) -> Vec { vec!["infer_mime_type".to_string()] } fn options(&self) -> &std::collections::HashMap { self.base.options() } fn options_mut( &mut self, ) -> anyhow::Result<&mut std::collections::HashMap> { Ok(self.base.options_mut()) } fn parallel_safe(&self) -> bool { true } } #[ctor::ctor] fn register_infer_plugin() { register_meta_plugin(MetaPluginType::Infer, |options, outputs| { Box::new(InferMetaPlugin::new(options, outputs)) }) .expect("Failed to register InferMetaPlugin"); }