Fixes #3952.
[librarian.git] / librarian / embeds / __init__.py
1 from __future__ import unicode_literals
2
3 import importlib
4 from lxml import etree
5
6 known_types = {
7     'application/mathml+xml': 'librarian.embeds.mathml.MathML',
8     'application/x-latex': 'librarian.embeds.latex.LaTeX',
9 }
10
11 class Embed():
12     @classmethod
13     def transforms_to(cls, mime_types, downgrade=False):
14         matches = set()
15         for name, method in cls.__dict__.iteritems():
16             if hasattr(method, "embed_converts_to"):
17                 conv_type, conv_downgrade = method.embed_converts_to
18                 if downgrade == conv_downgrade and conv_type in mime_types:
19                     matches.add(conv_type)
20         return matches
21
22     def transform_to(self, mime_type, downgrade=False):
23         for name, method in type(cls).__dict__.iteritems():
24             if hasattr(method, "embed_converts_to"):
25                 conv_type, conv_downgrade = method.embed_converts_to
26                 if downgrade == conv_downgrade and conv_type == mime_type:
27                     return method(self)
28
29
30 class DataEmbed(Embed):
31     def __init__(self, data=None):
32         self.data = data
33
34 class TreeEmbed(Embed):
35     def __init__(self, tree=None):
36         if isinstance(tree, etree._Element):
37             tree = etree.ElementTree(tree)
38         self.tree = tree
39
40 def converts_to(mime_type, downgrade=False):
41     def decorator(method):
42         method.embed_converts_to = mime_type, downgrade
43         return method
44     return decorator
45
46 def downgrades_to(mime_type):
47     return converts_to(mime_type, True)
48
49 def create_embed(mime_type, tree=None, data=None):
50     embed = known_types.get(mime_type)
51     if embed is None:
52         embed = DataEmbed if tree is None else TreeEmbed
53     else:
54         mod_name, cls_name = embed.rsplit('.', 1)
55         mod = importlib.import_module(mod_name)
56         embed = getattr(mod, cls_name)
57
58     return embed(data if tree is None else tree)