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("&amp;")
  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')