Add support for media files (#8)

Fixes #4 . Thanks to user otkoge (https://github.com/otkoge) for #6 , which this is based on.
kerrickstaley · May 7, 2017 · 8cbcc3c · 8cbcc3c
1 parent d0b73e6
commit 8cbcc3c
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -69,6 +69,21 @@ genanki.Package(my_deck).write_to_file('output.apkg')
 
 You can then load `output.apkg` into Anki using File -> Import...
 
+## Media Files
+To add sounds or images, set the `media_files` attribute on your `Package`:
+
+```python
+my_package = genanki.Package(my_deck)
+my_package.media_files = ['sound.mp3', 'image.jpg']
+```
+
+The media files should be in the current working directory. They can be referenced in notes like this:
+
+```html
+[sound:sound.mp3]
+<img src="image.jpg">
+```
+
 ## Note GUIDs
 `Note`s have a `guid` property that uniquely identifies the note. If you import a new note that has the same GUID as an
 existing note, the new note will overwrite the old one (as long as their models have the same fields).

diff --git a/genanki/__init__.py b/genanki/__init__.py
@@ -272,12 +272,14 @@ def write_to_db(self, cursor, now_ts):
 
 
 class Package:
-  def __init__(self, deck_or_decks=None):
+  def __init__(self, deck_or_decks=None, media_files=None):
     if isinstance(deck_or_decks, Deck):
       self.decks = [deck_or_decks]
     else:
       self.decks = deck_or_decks
 
+    self.media_files = media_files or []
+
   def write_to_file(self, file):
     dbfile, dbfilename = tempfile.mkstemp()
     os.close(dbfile)
@@ -293,7 +295,12 @@ def write_to_file(self, file):
 
     with zipfile.ZipFile(file, 'w') as outzip:
       outzip.write(dbfilename, 'collection.anki2')
-      outzip.writestr('media', '{}')
+
+      media_json = dict(enumerate(self.media_files))
+      outzip.writestr('media', json.dumps(media_json))
+
+      for i, f in media_json.items():
+        outzip.write(f, str(i))
 
   def write_to_db(self, cursor, now_ts):
     cursor.executescript(APKG_SCHEMA)

diff --git a/tests/test_genanki.py b/tests/test_genanki.py
@@ -64,6 +64,21 @@
   ],
 )
 
+# VALID_MP3 and VALID_JPG courtesy of https://github.com/mathiasbynens/small
+VALID_MP3 = (
+  b'\xff\xe3\x18\xc4\x00\x00\x00\x03H\x00\x00\x00\x00LAME3.98.2\x00\x00\x00'
+  b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
+  b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
+  b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00')
+
+VALID_JPG = (
+  b'\xff\xd8\xff\xdb\x00C\x00\x03\x02\x02\x02\x02\x02\x03\x02\x02\x02\x03\x03'
+  b'\x03\x03\x04\x06\x04\x04\x04\x04\x04\x08\x06\x06\x05\x06\t\x08\n\n\t\x08\t'
+  b'\t\n\x0c\x0f\x0c\n\x0b\x0e\x0b\t\t\r\x11\r\x0e\x0f\x10\x10\x11\x10\n\x0c'
+  b'\x12\x13\x12\x10\x13\x0f\x10\x10\x10\xff\xc9\x00\x0b\x08\x00\x01\x00\x01'
+  b'\x01\x01\x11\x00\xff\xcc\x00\x06\x00\x10\x10\x05\xff\xda\x00\x08\x01\x01'
+  b'\x00\x00?\x00\xd2\xcf \xff\xd9')
+
 
 class TestWithCollection:
   def setup(self):
@@ -172,3 +187,24 @@ def guid(self):
 
     # test passes if this doesn't raise an exception
     MyNote()
+
+  def test_media_files(self):
+    # change to a scratch directory so we can write files
+    os.chdir(tempfile.mkdtemp())
+
+    deck = genanki.Deck(123456, 'foodeck')
+    note = genanki.Note(TEST_MODEL, [
+      'question [sound:present.mp3] [sound:missing.mp3]',
+      'answer <img src="present.jpg"> <img src="missing.jpg">'])
+    deck.add_note(note)
+
+    # populate files with data
+    with open('present.mp3', 'wb') as h:
+      h.write(VALID_MP3)
+    with open('present.jpg', 'wb') as h:
+      h.write(VALID_JPG)
+
+    self.import_package(genanki.Package(deck, media_files=['present.mp3', 'present.jpg']))
+
+    missing, unused, invalid = self.col.media.check()
+    assert set(missing) == {'missing.mp3', 'missing.jpg'}