Fix #1002708 (FB2 metadata save: cover image update)

src/calibre/ebooks/metadata/fb2.py

@@ -5,11 +5,15 @@ __copyright__ = '2011, Roman Mukhin <ramses_ru at hotmail.com>, '\
                 '2008, Anatoly Shipitsin <norguhtar at gmail.com>'
 '''Read meta information from fb2 files'''
 
-import os
-import datetime
+import os, random, datetime
 from functools import partial
+from string import ascii_letters, digits
+from base64 import b64encode
+
 from lxml import etree
+
 from calibre.utils.date import parse_date
+from calibre.utils.magick.draw import save_cover_data_to
 from calibre import guess_type, guess_all_extensions, prints, force_unicode
 from calibre.ebooks.metadata import MetaInformation, check_isbn
 from calibre.ebooks.chardet import xml_to_unicode
@@ -25,6 +29,9 @@ tostring = partial(etree.tostring, method='text', encoding=unicode)
 def FB2(tag):
     return '{%s}%s'%(NAMESPACES['fb2'], tag)
 
+def XLINK(tag):
+    return '{%s}%s'%(NAMESPACES['xlink'], tag)
+
 def get_metadata(stream):
     ''' Return fb2 metadata as a L{MetaInformation} object '''
 
@@ -308,20 +315,48 @@ def _set_series(title_info, mi):
         except:
             seq.set('number', '1')
 
-def _create_tag(parent, tag, at_start=True):
+def _rnd_name(size=8, chars=ascii_letters + digits):
+    return ''.join(random.choice(chars) for x in range(size))
+
+def _rnd_pic_file_name(prefix='calibre_cover_', size=32, ext='jpg'):
+    return prefix + _rnd_name(size=size) + '.' + ext
+
+def _encode_into_jpeg(data):
+    data = save_cover_data_to(data, 'cover.jpg', return_data=True)
+    return b64encode(data)
+
+def _set_cover(title_info, mi):
+    if not mi.is_null('cover_data') and mi.cover_data[1]:
+        coverpage = _get_or_create(title_info, 'coverpage')
+        cim_tag = _get_or_create(coverpage, 'image')
+        if cim_tag.attrib.has_key(XLINK('href')):
+            cim_filename = cim_tag.attrib[XLINK('href')][1:]
+        else:
+            cim_filename = _rnd_pic_file_name('cover')
+            cim_tag.attrib[XLINK('href')] = '#' + cim_filename
+        fb2_root = cim_tag.getroottree().getroot()
+        cim_binary = _get_or_create(fb2_root, 'binary', attribs={'id': cim_filename}, at_start=False)
+        cim_binary.attrib['content-type'] = 'image/jpeg'
+        cim_binary.text = _encode_into_jpeg(mi.cover_data[1])
+
+def _create_tag(parent, tag, attribs={}, at_start=True):
     ans = parent.makeelement(FB2(tag))
+    ans.attrib.update(attribs)
     if at_start:
         parent.insert(0, ans)
     else:
         parent.append(ans)
     return ans
 
-def _get_or_create(parent, tag, at_start=True):
-    ans = XPath('./fb2:'+tag)(parent)
+def _get_or_create(parent, tag, attribs={}, at_start=True):
+    xpathstr='./fb2:'+tag
+    for n, v in attribs.items():
+        xpathstr += '[@%s="%s"]' % (n, v)
+    ans = XPath(xpathstr)(parent)
     if ans:
         ans = ans[0]
     else:
-        ans = _create_tag(parent, tag, at_start)
+        ans = _create_tag(parent, tag, attribs, at_start)
     return ans
 
 def set_metadata(stream, mi, apply_null=False, update_timestamp=False):
@@ -337,6 +372,7 @@ def set_metadata(stream, mi, apply_null=False, update_timestamp=False):
     _set_tags(ti, mi)
     _set_authors(ti, mi)
     _set_title(ti, mi)
+    _set_cover(ti, mi)
 
     for child in ti:
         child.tail = indent