From 4096d21c07fc4c635ae8d9020fa34e6f3313ec84 Mon Sep 17 00:00:00 2001 From: Damien Elmes Date: Tue, 28 Jan 2020 21:45:26 +1000 Subject: [PATCH] add add_file() and write_data() --- proto/backend.proto | 18 ++++++---- pylib/anki/media.py | 74 ++++++++++++++++----------------------- pylib/anki/rsbackend.py | 13 +++++-- pylib/anki/storage.py | 3 +- pylib/tests/test_media.py | 4 +-- qt/aqt/editor.py | 7 ++-- rslib/src/backend.rs | 26 +++++++++++--- rslib/src/err.rs | 12 +++++++ rspy/src/lib.rs | 4 +-- 9 files changed, 97 insertions(+), 64 deletions(-) diff --git a/proto/backend.proto b/proto/backend.proto index 8c240d5a3..c41a80386 100644 --- a/proto/backend.proto +++ b/proto/backend.proto @@ -19,6 +19,7 @@ message BackendInput { string strip_av_tags = 23; ExtractAVTagsIn extract_av_tags = 24; string expand_clozes_to_reveal_latex = 25; + AddFileToMediaFolderIn add_file_to_media_folder = 26; } } @@ -34,6 +35,7 @@ message BackendOutput { string strip_av_tags = 23; ExtractAVTagsOut extract_av_tags = 24; string expand_clozes_to_reveal_latex = 25; + string add_file_to_media_folder = 26; BackendError error = 2047; } @@ -41,16 +43,13 @@ message BackendOutput { message BackendError { oneof value { - InvalidInputError invalid_input = 1; - TemplateParseError template_parse = 2; + StringError invalid_input = 1; + StringError template_parse = 2; + StringError io_error = 3; } } -message InvalidInputError { - string info = 1; -} - -message TemplateParseError { +message StringError { string info = 1; bool q_side = 2; } @@ -174,3 +173,8 @@ message TTSTag { float speed = 4; repeated string other_args = 5; } + +message AddFileToMediaFolderIn { + string desired_name = 1; + bytes data = 2; +} \ No newline at end of file diff --git a/pylib/anki/media.py b/pylib/anki/media.py index 3cc54a87e..0549976ac 100644 --- a/pylib/anki/media.py +++ b/pylib/anki/media.py @@ -25,6 +25,10 @@ from anki.latex import render_latex from anki.utils import checksum, isMac, isWin +def media_folder_from_col_path(col_path: str) -> str: + return re.sub(r"(?i)\.(anki2)$", ".media", col_path) + + class MediaManager: soundRegexps = [r"(?i)(\[sound:(?P[^]]+)\])"] @@ -43,7 +47,7 @@ class MediaManager: self._dir = None return # media directory - self._dir = re.sub(r"(?i)\.(anki2)$", ".media", self.col.path) + self._dir = media_folder_from_col_path(self.col.path) if not os.path.exists(self._dir): os.makedirs(self._dir) try: @@ -155,58 +159,42 @@ create table meta (dirMod int, lastUsn int); insert into meta values (0, 0); # Adding media ########################################################################## - # opath must be in unicode - def addFile(self, opath: str) -> Any: - with open(opath, "rb") as f: - return self.writeData(opath, f.read()) + def add_file(self, path: str) -> str: + """Add basename of path to the media folder, renaming if not unique. - def writeData(self, opath: str, data: bytes, typeHint: Optional[str] = None) -> Any: - # if fname is a full path, use only the basename - fname = os.path.basename(opath) + Returns possibly-renamed filename.""" + with open(path, "rb") as f: + return self.write_data(os.path.basename(path), f.read()) - # if it's missing an extension and a type hint was provided, use that - if not os.path.splitext(fname)[1] and typeHint: + def write_data(self, desired_fname: str, data: bytes) -> str: + """Write the file to the media folder, renaming if not unique. + + Returns possibly-renamed filename.""" + return self.col.backend.add_file_to_media_folder(desired_fname, data) + + def add_extension_based_on_mime(self, fname: str, content_type: str) -> str: + "If jpg or png mime, add .png/.jpg if missing extension." + if not os.path.splitext(fname)[1]: # mimetypes is returning '.jpe' even after calling .init(), so we'll do # it manually instead - typeMap = { + type_map = { "image/jpeg": ".jpg", "image/png": ".png", } - if typeHint in typeMap: - fname += typeMap[typeHint] + if content_type in type_map: + fname += type_map[content_type] + return fname - # make sure we write it in NFC form (pre-APFS Macs will autoconvert to NFD), - # and return an NFC-encoded reference - fname = unicodedata.normalize("NFC", fname) - # ensure it's a valid filename - base = self.cleanFilename(fname) - (root, ext) = os.path.splitext(base) + # legacy + addFile = add_file - def repl(match): - n = int(match.group(1)) - return " (%d)" % (n + 1) - - # find the first available name - csum = checksum(data) - while True: - fname = root + ext - path = os.path.join(self.dir(), fname) - # if it doesn't exist, copy it directly - if not os.path.exists(path): - with open(path, "wb") as f: - f.write(data) - return fname - # if it's identical, reuse - with open(path, "rb") as f: - if checksum(f.read()) == csum: - return fname - # otherwise, increment the index in the filename - reg = r" \((\d+)\)$" - if not re.search(reg, root): - root = root + " (1)" - else: - root = re.sub(reg, repl, root) + # legacy + def writeData(self, opath: str, data: bytes, typeHint: Optional[str] = None) -> str: + fname = os.path.basename(opath) + if typeHint: + fname = self.add_extension_based_on_mime(fname, typeHint) + return self.write_data(fname, data) # String manipulation ########################################################################## diff --git a/pylib/anki/rsbackend.py b/pylib/anki/rsbackend.py index 253c31231..699e070eb 100644 --- a/pylib/anki/rsbackend.py +++ b/pylib/anki/rsbackend.py @@ -90,8 +90,8 @@ def proto_replacement_list_to_native( class RustBackend: - def __init__(self, path: str): - self._backend = ankirspy.Backend(path) + def __init__(self, col_path: str, media_folder: str): + self._backend = ankirspy.Backend(col_path, media_folder) def _run_command(self, input: pb.BackendInput) -> pb.BackendOutput: input_bytes = input.SerializeToString() @@ -181,3 +181,12 @@ class RustBackend: return self._run_command( pb.BackendInput(expand_clozes_to_reveal_latex=text) ).expand_clozes_to_reveal_latex + + def add_file_to_media_folder(self, desired_name: str, data: bytes) -> str: + return self._run_command( + pb.BackendInput( + add_file_to_media_folder=pb.AddFileToMediaFolderIn( + desired_name=desired_name, data=data + ) + ) + ).add_file_to_media_folder diff --git a/pylib/anki/storage.py b/pylib/anki/storage.py index 3950e77a6..4c4b54735 100644 --- a/pylib/anki/storage.py +++ b/pylib/anki/storage.py @@ -11,6 +11,7 @@ from anki.collection import _Collection from anki.consts import * from anki.db import DB from anki.lang import _ +from anki.media import media_folder_from_col_path from anki.rsbackend import RustBackend from anki.stdmodels import ( addBasicModel, @@ -30,7 +31,7 @@ def Collection( path: str, lock: bool = True, server: Optional[ServerData] = None, log: bool = False ) -> _Collection: "Open a new or existing collection. Path must be unicode." - backend = RustBackend(path) + backend = RustBackend(path, media_folder_from_col_path(path)) assert path.endswith(".anki2") path = os.path.abspath(path) create = not os.path.exists(path) diff --git a/pylib/tests/test_media.py b/pylib/tests/test_media.py index ca2fa41a4..1ff9b16d5 100644 --- a/pylib/tests/test_media.py +++ b/pylib/tests/test_media.py @@ -17,10 +17,10 @@ def test_add(): assert d.media.addFile(path) == "foo.jpg" # adding the same file again should not create a duplicate assert d.media.addFile(path) == "foo.jpg" - # but if it has a different md5, it should + # but if it has a different sha1, it should with open(path, "w") as f: f.write("world") - assert d.media.addFile(path) == "foo (1).jpg" + assert d.media.addFile(path) == "foo-7c211433f02071597741e6ff5a8ea34789abbf43.jpg" def test_strings(): diff --git a/qt/aqt/editor.py b/qt/aqt/editor.py index b5974e371..0ef34eb2b 100644 --- a/qt/aqt/editor.py +++ b/qt/aqt/editor.py @@ -793,8 +793,11 @@ to a cloze type first, via Edit>Change Note Type.""" self.mw.progress.finish() # strip off any query string url = re.sub(r"\?.*?$", "", url) - path = urllib.parse.unquote(url) - return self.mw.col.media.writeData(path, filecontents, typeHint=ct) + fname = os.path.basename(urllib.parse.unquote(url)) + if ct: + fname = self.mw.col.media.add_extension_based_on_mime(fname, ct) + + return self.mw.col.media.write_data(fname, filecontents) # Paste/drag&drop ###################################################################### diff --git a/rslib/src/backend.rs b/rslib/src/backend.rs index c0a318d7e..9fc70a24e 100644 --- a/rslib/src/backend.rs +++ b/rslib/src/backend.rs @@ -6,6 +6,7 @@ use crate::backend_proto::backend_input::Value; use crate::backend_proto::RenderedTemplateReplacement; use crate::cloze::expand_clozes_to_reveal_latex; use crate::err::{AnkiError, Result}; +use crate::media::add_data_to_folder_uniquely; use crate::sched::{local_minutes_west_for_stamp, sched_timing_today}; use crate::template::{ render_card, without_legacy_template_directives, FieldMap, FieldRequirements, ParsedTemplate, @@ -18,7 +19,8 @@ use std::path::PathBuf; pub struct Backend { #[allow(dead_code)] - path: PathBuf, + col_path: PathBuf, + media_folder: PathBuf, } /// Convert an Anki error to a protobuf error. @@ -26,10 +28,11 @@ impl std::convert::From for pt::BackendError { fn from(err: AnkiError) -> Self { use pt::backend_error::Value as V; let value = match err { - AnkiError::InvalidInput { info } => V::InvalidInput(pt::InvalidInputError { info }), + AnkiError::InvalidInput { info } => V::InvalidInput(pt::StringError { info }), AnkiError::TemplateError { info, q_side } => { V::TemplateParse(pt::TemplateParseError { info, q_side }) - } + }, + AnkiError::IOError { info } => V::IoError(pt::StringError { info }), }; pt::BackendError { value: Some(value) } @@ -44,8 +47,11 @@ impl std::convert::From for pt::backend_output::Value { } impl Backend { - pub fn new>(path: P) -> Backend { - Backend { path: path.into() } + pub fn new>(col_path: P, media_folder: P) -> Backend { + Backend { + col_path: col_path.into(), + media_folder: media_folder.into(), + } } /// Decode a request, process it, and return the encoded result. @@ -107,6 +113,9 @@ impl Backend { Value::ExpandClozesToRevealLatex(input) => { OValue::ExpandClozesToRevealLatex(expand_clozes_to_reveal_latex(&input)) } + Value::AddFileToMediaFolder(input) => { + OValue::AddFileToMediaFolder(self.add_file_to_media_folder(input)?) + } }) } @@ -219,6 +228,13 @@ impl Backend { av_tags: pt_tags, } } + + fn add_file_to_media_folder(&self, input: pt::AddFileToMediaFolderIn) -> Result { + Ok( + add_data_to_folder_uniquely(&self.media_folder, &input.desired_name, &input.data)? + .into(), + ) + } } fn ords_hash_to_set(ords: HashSet) -> Vec { diff --git a/rslib/src/err.rs b/rslib/src/err.rs index ceb9626bf..7f3fd4c2e 100644 --- a/rslib/src/err.rs +++ b/rslib/src/err.rs @@ -2,6 +2,7 @@ // License: GNU AGPL, version 3 or later; http://www.gnu.org/licenses/agpl.html pub use failure::{Error, Fail}; +use std::io; pub type Result = std::result::Result; @@ -12,6 +13,9 @@ pub enum AnkiError { #[fail(display = "invalid card template: {}", info)] TemplateError { info: String, q_side: bool }, + + #[fail(display = "I/O error: {}", info)] + IOError { info: String }, } // error helpers @@ -34,3 +38,11 @@ pub enum TemplateError { field: String, }, } + +impl From for AnkiError { + fn from(err: io::Error) -> Self { + AnkiError::IOError { + info: format!("{:?}", err), + } + } +} diff --git a/rspy/src/lib.rs b/rspy/src/lib.rs index d7eaf7a4f..b03b1fd6a 100644 --- a/rspy/src/lib.rs +++ b/rspy/src/lib.rs @@ -16,10 +16,10 @@ fn buildhash() -> &'static str { #[pymethods] impl Backend { #[new] - fn init(obj: &PyRawObject, path: String) { + fn init(obj: &PyRawObject, col_path: String, media_folder: String) { obj.init({ Backend { - backend: RustBackend::new(path), + backend: RustBackend::new(col_path, media_folder), } }); }