Add @Format directive for '.'-separated formatted resources

alekstorm · Sep 16, 2011 · 3df52ca · 3df52ca
1 parent d557375
commit 3df52ca
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 2 deletions.
diff --git a/test/test.py b/test/test.py
@@ -56,12 +56,18 @@ class LogoutResource(AppResource):
     def get(self, request, user):
         return HTTPResponse(entity='Bye, '+user, cookies={'user': ''})
 
+class FormatResource(Resource):
+    @format({'json': lambda r: JSONResource(r), 'xml': json2xml})
+    def __getitem__(self, name):
+        return {'name': 'Guido', 'interests': [{'title': 'python', 'level': 9}, {'title': 'dancing', 'level': 5}], 'books': ['Python Tutorial', 'Python Reference Manual'], 'address': {'city': 'Mountain View', 'state': 'CA'}}
+
 app = Application(Memcacher(Client(['127.0.0.1:11211']), {
     '': lambda request: 'Hello World!',
     'static': StaticDirectoryResource(os.path.join(os.path.dirname(__file__), 'static')),
     'json': JSONResource({'a': ['b', 1], 'c': {'d': True}}),
     'args': ArgResource(),
     'memcached': MemcachedResource(),
+    'format': FormatResource(),
     'auth': {
         'login': LoginResource(),
         'secret': SecretResource(),

diff --git a/vortex/app.py b/vortex/app.py
@@ -2,14 +2,55 @@
 import httplib
 import json
 import logging
+try:
+    import cStringIO as StringIO
+except ImportError:
+    import StringIO
 import tornado.web
 import traceback
 import urllib
+from   xml.etree.ElementTree import Element, ElementTree, iselement
 
 from vortex.response import *
 
 logger = logging.getLogger('vortex')
 
+def json2xml(data):
+    def convert_elem(data, tag='item'):
+        root = Element(tag)
+        if isinstance(data, dict):
+            for key, val in data.iteritems():
+                if isinstance(val, dict) or isinstance(val, list):
+                    root.append(convert_elem(val, key))
+                else:
+                    root.set(key, str(val))
+        elif isinstance(data, list):
+            for item in data:
+                root.append(convert_elem(item))
+        else:
+            root.set('_value', str(data))
+        return root
+
+    return lambda request: convert_elem(data, 'root')
+
+def format(handlers, default=None, unknown=None):
+    def wrap1(getitem):
+        def wrap2(self, name):
+            parts = name.split('.')
+            resource = getitem(self, parts[0])
+            if len(parts) > 1:
+                handler = handlers.get(parts[1], None)
+                if handler:
+                    return handler(resource)
+                if unknown:
+                    return unknown(resource)
+                raise KeyError()
+            if default:
+                return default(resource)
+            raise KeyError()
+        return wrap2
+    return wrap1
+
 def authenticate(retrieve, cookie_name, redirect=None, unauthorized=None):
     def wrap1(fn):
         def wrap2(self, request, *args, **kwargs):
@@ -57,8 +98,11 @@ def coerce_response(response):
     elif isinstance(response, basestring):
         response = HTTPResponse(entity=response)
     elif isinstance(response, dict):
-        response = HTTPResponse(entity=json.dumps(response))
-        response.headers.setdefault('Content-Type', 'application/json')
+        response = HTTPResponse(entity=json.dumps(response), headers={'Content-Type': 'application/json'})
+    elif iselement(response):
+        xml = StringIO.StringIO()
+        ElementTree(response).write(xml)
+        response = HTTPResponse(entity=xml.getvalue(), headers={'Content-Type': 'application/xml'})
     return response