Merge branch 'master' into python3
[wrw.git] / wrw / sp / util.py
1 from . import cons
2
3 def findnsnames(el):
4     names = {}
5     nid = [1]
6     def proc(el):
7         if isinstance(el, cons.element):
8             if el.ns not in names:
9                 names[el.ns] = "n" + str(nid[0])
10                 nid[:] = [nid[0] + 1]
11             for ch in el.children:
12                 proc(ch)
13     proc(el)
14     if None in names:
15         names[None] = None
16     else:
17         names[el.ns] = None
18     return names
19
20 class formatter(object):
21     def __init__(self, out, root, nsnames=None, charset="utf-8", doctype=None):
22         self.root = root
23         if nsnames is None:
24             nsnames = findnsnames(root)
25         self.nsnames = nsnames
26         self.out = out
27         self.charset = charset
28         self.doctype = doctype
29
30     def write(self, text):
31         self.out.write(text.encode(self.charset))
32
33     def quotewrite(self, buf):
34         for ch in buf:
35             if ch == '&':
36                 self.write("&")
37             elif ch == '<':
38                 self.write("&lt;")
39             elif ch == '>':
40                 self.write("&gt;")
41             else:
42                 self.write(ch)
43
44     def text(self, el):
45         self.quotewrite(el)
46
47     def rawcode(self, el):
48         self.write(el)
49
50     def attrval(self, buf):
51         qc, qt = ("'", "&apos;") if '"' in buf else ('"', "&quot;")
52         self.write(qc)
53         for ch in buf:
54             if ch == '&':
55                 self.write("&amp;")
56             elif ch == '<':
57                 self.write("&lt;")
58             elif ch == '>':
59                 self.write("&gt;")
60             elif ch == qc:
61                 self.write(qt)
62             else:
63                 self.write(ch)
64         self.write(qc)
65
66     def attr(self, k, v):
67         self.write(k)
68         self.write('=')
69         self.attrval(v)
70
71     def shorttag(self, el, **extra):
72         self.write('<' + self.elname(el))
73         for k, v in el.attrs.items():
74             self.write(' ')
75             self.attr(k, v)
76         for k, v in extra.items():
77             self.write(' ')
78             self.attr(k, v)
79         self.write(" />")
80
81     def elname(self, el):
82         ns = self.nsnames[el.ns]
83         if ns is None:
84             return el.name
85         else:
86             return ns + ':' + el.name
87
88     def starttag(self, el, **extra):
89         self.write('<' + self.elname(el))
90         for k, v in el.attrs.items():
91             self.write(' ')
92             self.attr(k, v)
93         for k, v in extra.items():
94             self.write(' ')
95             self.attr(k, v)
96         self.write('>')
97
98     def endtag(self, el):
99         self.write('</' + self.elname(el) + '>')
100
101     def longtag(self, el):
102         self.starttag(el, **extra)
103         for ch in el.children:
104             self.node(ch)
105         self.endtag(el)
106
107     def element(self, el, **extra):
108         if len(el.children) == 0:
109             self.shorttag(el, **extra)
110         else:
111             self.longtag(el, **extra)
112
113     def node(self, el):
114         if isinstance(el, cons.element):
115             self.element(el)
116         elif isinstance(el, cons.text):
117             self.text(el)
118         elif isinstance(el, cons.raw):
119             self.rawcode(el)
120         else:
121             raise Exception("Unknown object in element tree: " + el)
122
123     def start(self):
124         self.write('<?xml version="1.0" encoding="' + self.charset + '" ?>\n')
125         if self.doctype:
126             self.write('<!DOCTYPE %s PUBLIC "%s" "%s">\n' % (self.root.name,
127                                                              self.doctype[0],
128                                                              self.doctype[1]))
129         extra = {}
130         for uri, nm in self.nsnames.items():
131             if uri is None:
132                 continue
133             if nm is None:
134                 extra["xmlns"] = uri
135             else:
136                 extra["xmlns:" + nm] = uri
137         self.element(self.root, **extra)
138
139     @classmethod
140     def output(cls, out, el, *args, **kw):
141         cls(out=out, root=el, *args, **kw).start()
142
143     @classmethod
144     def fragment(cls, out, el, *args, **kw):
145         cls(out=out, root=el, *args, **kw).node(el)
146
147     def update(self, **ch):
148         ret = type(self).__new__(type(self))
149         ret.__dict__.update(self.__dict__)
150         ret.__dict__.update(ch)
151         return ret
152
153 class iwriter(object):
154     def __init__(self, out):
155         self.out = out
156         self.atbol = True
157         self.col = 0
158
159     def write(self, buf):
160         for i in range(len(buf)):
161             c = buf[i:i + 1]
162             if c == b'\n':
163                 self.col = 0
164             else:
165                 self.col += 1
166             self.out.write(c)
167         self.atbol = False
168
169     def indent(self, indent):
170         if self.atbol:
171             return
172         if self.col != 0:
173             self.write(b'\n')
174         self.write(indent)
175         self.atbol = True
176
177 class indenter(formatter):
178     def __init__(self, indent="  ", *args, **kw):
179         super(indenter, self).__init__(*args, **kw)
180         self.out = iwriter(self.out)
181         self.indent = indent
182         self.curind = ""
183
184     def simple(self, el):
185         for ch in el.children:
186             if not isinstance(ch, cons.text):
187                 return False
188         return True
189
190     def longtag(self, el, **extra):
191         self.starttag(el, **extra)
192         sub = self
193         reind = False
194         if not self.simple(el):
195             sub = self.update(curind=self.curind + self.indent)
196             sub.reindent()
197             reind = True
198         for ch in el.children:
199             sub.node(ch)
200         if reind:
201             self.reindent()
202         self.endtag(el)
203
204     def element(self, el, **extra):
205         super(indenter, self).element(el, **extra)
206         if self.out.col > 80 and self.simple(el):
207             self.reindent()
208
209     def reindent(self):
210         self.out.indent(self.curind.encode(self.charset))
211
212     def start(self):
213         super(indenter, self).start()
214         self.write('\n')