add add_file() and write_data()

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;
+}

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<fname>[^]]+)\])"]
@@ -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
     ##########################################################################

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

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)

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():

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
     ######################################################################

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<AnkiError> 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<AnkiError> for pt::backend_output::Value {
 }
 
 impl Backend {
-    pub fn new<P: Into<PathBuf>>(path: P) -> Backend {
-        Backend { path: path.into() }
+    pub fn new<P: Into<PathBuf>>(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<String> {
+        Ok(
+            add_data_to_folder_uniquely(&self.media_folder, &input.desired_name, &input.data)?
+                .into(),
+        )
+    }
 }
 
 fn ords_hash_to_set(ords: HashSet<u16>) -> Vec<u32> {

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<T> = std::result::Result<T, AnkiError>;
 
@@ -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<io::Error> for AnkiError {
+    fn from(err: io::Error) -> Self {
+        AnkiError::IOError {
+            info: format!("{:?}", err),
+        }
+    }
+}

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