Factor out media importing

This commit is contained in:
RumovZ 2022-04-11 08:44:18 +02:00
parent fe878dcdc8
commit 251bfc4920
3 changed files with 118 additions and 97 deletions

View file

@ -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<String, (bool, SafeMediaEntry)>);
impl<'a> Context<'a> {
pub(super) fn prepare_media(&mut self) -> Result<MediaUseMap> {
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<HashMap<String, [u8; 20]>> {
let mgr = MediaManager::new(&self.media_folder, &self.media_db)?;
mgr.all_checksums(|_| true, &self.log)
}
}
fn prepare_media(
archive: &mut ZipArchive<File>,
existing_sha1s: &HashMap<String, [u8; 20]>,
) -> Result<MediaUseMap> {
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<String>, 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<Item = &SafeMediaEntry> {
self.0.values().filter_map(|t| t.0.then(|| &t.1))
}
}
impl SafeMediaEntry {
fn with_hash_from_archive(&mut self, archive: &mut ZipArchive<File>) -> 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)
}
}

View file

@ -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<String, (bool, SafeMediaEntry)>);
struct Context<'a> {
target_col: &'a mut Collection,
archive: ZipArchive<File>,
@ -42,39 +26,6 @@ struct Context<'a> {
usn: Usn,
}
impl MediaUseMap {
fn add(&mut self, filename: impl Into<String>, 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<Item = &SafeMediaEntry> {
self.0.values().filter_map(|t| t.0.then(|| &t.1))
}
}
impl SafeMediaEntry {
fn with_hash_from_archive(&mut self, archive: &mut ZipArchive<File>) -> 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<Path>) -> Result<OpOutput<()>> {
let file = File::open(path)?;
@ -85,29 +36,6 @@ impl Collection {
ctx.import()
})
}
fn all_existing_sha1s(&mut self) -> Result<HashMap<String, [u8; 20]>> {
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<File>,
search: impl TryIntoSearch,
with_scheduling: bool,
) -> Result<Self> {
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<MediaUseMap> {
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<File>,
search: impl TryIntoSearch,
with_scheduling: bool,
) -> Result<Self> {
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)
}
}

View file

@ -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,
};