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