From 251bfc492019377515df7e5cd063061f5e4f91df Mon Sep 17 00:00:00 2001 From: RumovZ Date: Mon, 11 Apr 2022 08:44:18 +0200 Subject: [PATCH] Factor out media importing --- .../package/apkg/import/media.rs | 99 +++++++++++++++ .../import_export/package/apkg/import/mod.rs | 114 +++--------------- .../package/apkg/import/notes.rs | 2 +- 3 files changed, 118 insertions(+), 97 deletions(-) create mode 100644 rslib/src/import_export/package/apkg/import/media.rs diff --git a/rslib/src/import_export/package/apkg/import/media.rs b/rslib/src/import_export/package/apkg/import/media.rs new file mode 100644 index 000000000..25ca2c5e1 --- /dev/null +++ b/rslib/src/import_export/package/apkg/import/media.rs @@ -0,0 +1,99 @@ +// Copyright: Ankitects Pty Ltd and contributors +// License: GNU AGPL, version 3 or later; http://www.gnu.org/licenses/agpl.html + +use std::{collections::HashMap, fs::File, mem}; + +use zip::ZipArchive; + +use super::Context; +use crate::{ + import_export::package::{ + media::{extract_media_entries, SafeMediaEntry}, + Meta, + }, + media::{ + files::{add_hash_suffix_to_file_stem, sha1_of_reader}, + MediaManager, + }, + prelude::*, +}; + +/// Map of source media files, that do not already exist in the target. +/// +/// original, normalized filename → (refererenced on import material, +/// entry with possibly remapped filename) +#[derive(Default)] +pub(super) struct MediaUseMap(HashMap); + +impl<'a> Context<'a> { + pub(super) fn prepare_media(&mut self) -> Result { + let existing_sha1s = self.target_col.all_existing_sha1s()?; + prepare_media(&mut self.archive, &existing_sha1s) + } + + pub(super) fn copy_media(&mut self, media_map: &mut MediaUseMap) -> Result<()> { + for entry in media_map.used_entries() { + entry.copy_from_archive(&mut self.archive, &self.target_col.media_folder)?; + } + Ok(()) + } +} + +impl Collection { + fn all_existing_sha1s(&mut self) -> Result> { + let mgr = MediaManager::new(&self.media_folder, &self.media_db)?; + mgr.all_checksums(|_| true, &self.log) + } +} + +fn prepare_media( + archive: &mut ZipArchive, + existing_sha1s: &HashMap, +) -> Result { + let mut media_map = MediaUseMap::default(); + for mut entry in extract_media_entries(&Meta::new_legacy(), archive)? { + if let Some(other_sha1) = existing_sha1s.get(&entry.name) { + entry.with_hash_from_archive(archive)?; + if entry.sha1 != *other_sha1 { + let original_name = entry.uniquify_name(); + media_map.add(original_name, entry); + } + } else { + media_map.add(entry.name.clone(), entry); + } + } + Ok(media_map) +} + +impl MediaUseMap { + pub(super) fn add(&mut self, filename: impl Into, entry: SafeMediaEntry) { + self.0.insert(filename.into(), (false, entry)); + } + + pub(super) fn use_entry(&mut self, filename: &str) -> Option<&SafeMediaEntry> { + self.0.get_mut(filename).map(|(used, entry)| { + *used = true; + &*entry + }) + } + + pub(super) fn used_entries(&self) -> impl Iterator { + self.0.values().filter_map(|t| t.0.then(|| &t.1)) + } +} + +impl SafeMediaEntry { + fn with_hash_from_archive(&mut self, archive: &mut ZipArchive) -> Result<()> { + if self.sha1 == [0; 20] { + let mut reader = self.fetch_file(archive)?; + self.sha1 = sha1_of_reader(&mut reader)?; + } + Ok(()) + } + + /// Requires sha1 to be set. Returns old file name. + fn uniquify_name(&mut self) -> String { + let new_name = add_hash_suffix_to_file_stem(&self.name, &self.sha1); + mem::replace(&mut self.name, new_name) + } +} diff --git a/rslib/src/import_export/package/apkg/import/mod.rs b/rslib/src/import_export/package/apkg/import/mod.rs index aace86345..a06752c5c 100644 --- a/rslib/src/import_export/package/apkg/import/mod.rs +++ b/rslib/src/import_export/package/apkg/import/mod.rs @@ -3,9 +3,10 @@ mod cards; mod decks; +mod media; mod notes; -use std::{collections::HashMap, fs::File, io, mem, path::Path}; +use std::{fs::File, io, path::Path}; pub(crate) use notes::NoteMeta; use tempfile::NamedTempFile; @@ -13,28 +14,11 @@ use zip::ZipArchive; use crate::{ collection::CollectionBuilder, - import_export::{ - gather::ExchangeData, - package::{ - media::{extract_media_entries, SafeMediaEntry}, - Meta, - }, - }, - media::{ - files::{add_hash_suffix_to_file_stem, sha1_of_reader}, - MediaManager, - }, + import_export::{gather::ExchangeData, package::Meta}, prelude::*, search::SearchNode, }; -/// Map of source media files, that do not already exist in the target. -/// -/// original, normalized filename → (refererenced on import material, -/// entry with possibly remapped filename) -#[derive(Default)] -struct MediaUseMap(HashMap); - struct Context<'a> { target_col: &'a mut Collection, archive: ZipArchive, @@ -42,39 +26,6 @@ struct Context<'a> { usn: Usn, } -impl MediaUseMap { - fn add(&mut self, filename: impl Into, entry: SafeMediaEntry) { - self.0.insert(filename.into(), (false, entry)); - } - - fn use_entry(&mut self, filename: &str) -> Option<&SafeMediaEntry> { - self.0.get_mut(filename).map(|(used, entry)| { - *used = true; - &*entry - }) - } - - fn used_entries(&self) -> impl Iterator { - self.0.values().filter_map(|t| t.0.then(|| &t.1)) - } -} - -impl SafeMediaEntry { - fn with_hash_from_archive(&mut self, archive: &mut ZipArchive) -> Result<()> { - if self.sha1 == [0; 20] { - let mut reader = self.fetch_file(archive)?; - self.sha1 = sha1_of_reader(&mut reader)?; - } - Ok(()) - } - - /// Requires sha1 to be set. Returns old file name. - fn uniquify_name(&mut self) -> String { - let new_name = add_hash_suffix_to_file_stem(&self.name, &self.sha1); - mem::replace(&mut self.name, new_name) - } -} - impl Collection { pub fn import_apkg(&mut self, path: impl AsRef) -> Result> { let file = File::open(path)?; @@ -85,29 +36,6 @@ impl Collection { ctx.import() }) } - - fn all_existing_sha1s(&mut self) -> Result> { - let mgr = MediaManager::new(&self.media_folder, &self.media_db)?; - mgr.all_checksums(|_| true, &self.log) - } -} - -impl ExchangeData { - fn gather_from_archive( - archive: &mut ZipArchive, - search: impl TryIntoSearch, - with_scheduling: bool, - ) -> Result { - let mut zip_file = archive.by_name(Meta::new_legacy().collection_filename())?; - let mut tempfile = NamedTempFile::new()?; - io::copy(&mut zip_file, &mut tempfile)?; - let mut col = CollectionBuilder::new(tempfile.path()).build()?; - - let mut data = ExchangeData::default(); - data.gather_data(&mut col, search, with_scheduling)?; - - Ok(data) - } } impl<'a> Context<'a> { @@ -130,28 +58,22 @@ impl<'a> Context<'a> { self.import_cards_and_revlog(&imported_notes, &imported_decks)?; self.copy_media(&mut media_map) } +} - fn prepare_media(&mut self) -> Result { - let mut media_map = MediaUseMap::default(); - let existing_sha1s = self.target_col.all_existing_sha1s()?; - for mut entry in extract_media_entries(&Meta::new_legacy(), &mut self.archive)? { - if let Some(other_sha1) = existing_sha1s.get(&entry.name) { - entry.with_hash_from_archive(&mut self.archive)?; - if entry.sha1 != *other_sha1 { - let original_name = entry.uniquify_name(); - media_map.add(original_name, entry); - } - } else { - media_map.add(entry.name.clone(), entry); - } - } - Ok(media_map) - } +impl ExchangeData { + fn gather_from_archive( + archive: &mut ZipArchive, + search: impl TryIntoSearch, + with_scheduling: bool, + ) -> Result { + let mut zip_file = archive.by_name(Meta::new_legacy().collection_filename())?; + let mut tempfile = NamedTempFile::new()?; + io::copy(&mut zip_file, &mut tempfile)?; + let mut col = CollectionBuilder::new(tempfile.path()).build()?; - fn copy_media(&mut self, media_map: &mut MediaUseMap) -> Result<()> { - for entry in media_map.used_entries() { - entry.copy_from_archive(&mut self.archive, &self.target_col.media_folder)?; - } - Ok(()) + let mut data = ExchangeData::default(); + data.gather_data(&mut col, search, with_scheduling)?; + + Ok(data) } } diff --git a/rslib/src/import_export/package/apkg/import/notes.rs b/rslib/src/import_export/package/apkg/import/notes.rs index a656a45b0..6a8c4b990 100644 --- a/rslib/src/import_export/package/apkg/import/notes.rs +++ b/rslib/src/import_export/package/apkg/import/notes.rs @@ -10,7 +10,7 @@ use std::{ use sha1::Sha1; -use super::{Context, MediaUseMap}; +use super::{media::MediaUseMap, Context}; use crate::{ import_export::package::media::safe_normalized_file_name, prelude::*, text::replace_media_refs, };