Add Content-Length to SP responses.
[wrw.git] / wrw / sp / util.py
index 96e9880..e15b947 100644 (file)
@@ -1,3 +1,5 @@
+import StringIO
+from wrw import dispatch
 import cons
 
 def findnsnames(el):
@@ -44,6 +46,9 @@ class formatter(object):
     def text(self, el):
         self.quotewrite(el)
 
+    def rawcode(self, el):
+        self.write(el)
+
     def attrval(self, buf):
         qc, qt = (u"'", u"'") if u'"' in buf else (u'"', u""")
         self.write(qc)
@@ -95,7 +100,7 @@ class formatter(object):
     def endtag(self, el):
         self.write(u'</' + self.elname(el) + u'>')
 
-    def longtag(self, el):
+    def longtag(self, el, **extra):
         self.starttag(el, **extra)
         for ch in el.children:
             self.node(ch)
@@ -112,6 +117,8 @@ class formatter(object):
             self.element(el)
         elif isinstance(el, cons.text):
             self.text(el)
+        elif isinstance(el, cons.raw):
+            self.rawcode(el)
         else:
             raise Exception("Unknown object in element tree: " + el)
 
@@ -135,6 +142,16 @@ class formatter(object):
     def output(cls, out, el, *args, **kw):
         cls(out=out, root=el, *args, **kw).start()
 
+    @classmethod
+    def fragment(cls, out, el, *args, **kw):
+        cls(out=out, root=el, *args, **kw).node(el)
+
+    @classmethod
+    def format(cls, el, *args, **kw):
+        buf = StringIO.StringIO()
+        cls.output(buf, el, *args, **kw)
+        return buf.getvalue()
+
     def update(self, **ch):
         ret = type(self).__new__(type(self))
         ret.__dict__.update(self.__dict__)
@@ -183,19 +200,41 @@ class indenter(formatter):
         reind = False
         if not self.simple(el):
             sub = self.update(curind=self.curind + self.indent)
-            sub.out.indent(sub.curind)
+            sub.reindent()
             reind = True
         for ch in el.children:
             sub.node(ch)
         if reind:
-            self.out.indent(self.curind)
+            self.reindent()
         self.endtag(el)
 
     def element(self, el, **extra):
         super(indenter, self).element(el, **extra)
         if self.out.col > 80 and self.simple(el):
-            self.out.indent(self.curind)
+            self.reindent()
+
+    def reindent(self):
+        self.out.indent(self.curind.encode(self.charset))
 
     def start(self):
         super(indenter, self).start()
         self.write('\n')
+
+class response(dispatch.restart):
+    charset = "utf-8"
+    doctype = None
+    formatter = indenter
+
+    def __init__(self, root):
+        super(response, self).__init__()
+        self.root = root
+
+    @property
+    def ctype(self):
+        raise Exception("a subclass of wrw.sp.util.response must override ctype")
+
+    def handle(self, req):
+        ret = self.formatter.format(self.root, doctype=self.doctype, charset=self.charset)
+        req.ohead["Content-Type"] = self.ctype
+        req.ohead["Content-Length"] = len(ret)
+        return [ret]