obvious crash avoidance fixes
[librarian.git] / librarian / xmlutils.py
1 # -*- coding: utf-8 -*-
2 #
3 # This file is part of Librarian, licensed under GNU Affero GPLv3 or later.
4 # Copyright © Fundacja Nowoczesna Polska. See NOTICE for more information.
5 #
6 from lxml import etree
7 from collections import defaultdict
8
9
10 class Xmill(object):
11     """Transforms XML to some text.
12     Used instead of XSLT which is difficult and cumbersome.
13
14     """
15     def __init__(self, options=None):
16         self._options = []
17         if options:
18             self._options.append(options)
19         self.text_filters = []
20
21     def register_text_filter(self, fun):
22         self.text_filters.append(fun)
23
24     def filter_text(self, text):
25         for flt in self.text_filters:
26             if text is None:
27                 return None
28             text = flt(text)
29         return text
30
31     def generate(self, document):
32         """Generate text from node using handlers defined in class."""
33         output = self._handle_element(document)
34         return u''.join([x for x in flatten(output) if x is not None])
35
36     @property
37     def options(self):
38         """Returnes merged scoped options for current node.
39         """
40         # Here we can see how a decision not to return the modified map
41         # leads to a need for a hack.
42         return reduce(lambda a, b: a.update(b) or a, self._options, defaultdict(lambda: None))
43
44     @options.setter
45     def options(self, opts):
46         """Sets options overrides for current and child nodes
47         """
48         self._options.append(opts)
49
50
51     def _handle_for_element(self, element):
52         ns = None
53         tagname = None
54 #        from nose.tools import set_trace
55
56         if element.tag[0] == '{':
57             for nshort, nhref in element.nsmap.items():
58                 try:
59                     if element.tag.index('{%s}' % nhref) == 0:
60                         ns = nshort
61                         tagname  = element.tag[len('{%s}' % nhref):]
62                         break
63                 except ValueError:
64                     pass
65             if not ns:
66                 raise ValueError("Strange ns for tag: %s, nsmap: %s" %
67                                  (element.tag, element.nsmap))
68         else:
69             tagname = element.tag
70
71         if ns:
72             meth_name = "handle_%s__%s" % (ns, tagname)
73         else:
74             meth_name = "handle_%s" % (tagname,)
75
76         handler = getattr(self, meth_name, None)
77         return handler
78
79     def next(self, element):
80         if len(element):
81             return element[0]
82
83         while True:
84             sibling = element.getnext()
85             if sibling is not None: return sibling  # found a new branch to dig into
86             element = element.getparent()
87             if element is None: return None  # end of tree
88
89     def _handle_element(self, element):
90         if isinstance(element, etree._Comment): return None
91         
92         handler = self._handle_for_element(element)
93         # How many scopes
94         try:
95             options_scopes = len(self._options)
96
97             if handler is None:
98                 pre = [self.filter_text(element.text)]
99                 post = [self.filter_text(element.tail)]
100             else:
101                 vals = handler(element)
102                 # depending on number of returned values, vals can be None, a value, or a tuple.
103                 # how poorly designed is that? 9 lines below are needed just to unpack this.
104                 if vals is None:
105                     return [self.filter_text(element.tail)]
106                 else:
107                     if not isinstance(vals, tuple):
108                         return [vals, self.filter_text(element.tail)]
109                     else:
110                         pre = [vals[0], self.filter_text(element.text)]
111                         post = [vals[1], self.filter_text(element.tail)]
112
113             out = pre + [self._handle_element(child) for child in element] + post
114         finally:
115             # clean up option scopes if necessary
116             self._options = self._options[0:options_scopes]
117         return out
118
119
120 def tag_open_close(name_, classes_=None, **attrs):
121     u"""Creates tag beginning and end.
122     
123     >>> tag_open_close("a", "klass", x=u"ą<")
124     (u'<a x="\\u0105&lt;" class="klass">', u'</a>')
125
126     """
127     if classes_:
128         if isinstance(classes_, (tuple, list)): classes_ = ' '.join(classes_)
129         attrs['class'] = classes_
130
131     e = etree.Element(name_)
132     e.text = " "
133     for k, v in attrs.items():
134         e.attrib[k] = v
135     pre, post = etree.tostring(e, encoding=unicode).split(u"> <")
136     return pre + u">", u"<" + post
137
138 def tag(name_, classes_=None, **attrs):
139     """Returns a handler which wraps node contents in tag `name', with class attribute
140     set to `classes' and other attributes according to keyword paramters
141     """
142     def _hnd(self, element):
143         return tag_open_close(name_, classes_, **attrs)
144     return _hnd
145
146
147 def tagged(name, classes=None, **attrs):
148     """Handler decorator which wraps handler output in tag `name', with class attribute
149     set to `classes' and other attributes according to keyword paramters
150     """
151     if classes:
152         if isinstance(classes, (tuple,list)): classes = ' '.join(classes)
153         attrs['class'] = classes
154     a = ''.join([' %s="%s"' % (k,v) for (k,v) in attrs.items()])
155     def _decor(f):
156         def _wrap(self, element):
157             r = f(self, element)
158             if r is None: return
159
160             prepend = "<%s%s>" % (name, a)
161             append = "</%s>" % name
162
163             if isinstance(r, tuple):
164                 return prepend + r[0], r[1] + append
165             return prepend + r + append
166         return _wrap
167     return _decor
168
169
170 def ifoption(**options):
171     """Decorator which enables node only when options are set
172     """
173     def _decor(f):
174         def _handler(self, *args, **kw):
175             opts = self.options
176             for k, v in options.items():
177                 if opts[k] != v:
178                     return
179             return f(self, *args, **kw)
180         return _handler
181     return _decor
182
183 def flatten(l, ltypes=(list, tuple)):
184     """flatten function from BasicPropery/BasicTypes package
185     """
186     ltype = type(l)
187     l = list(l)
188     i = 0
189     while i < len(l):
190         while isinstance(l[i], ltypes):
191             if not l[i]:
192                 l.pop(i)
193                 i -= 1
194                 break
195             else:
196                 l[i:i + 1] = l[i]
197         i += 1
198     return ltype(l)