Commits

Geert Jansen committed ccb685d

Factor out PyYAML setup so that it can be reused by clients.

Comments (0)

Files changed (2)

lib/rest/entity/yaml.py

 
 import yaml
 from yaml import YAMLError
-from StringIO import StringIO
 
+from rest import http
+from rest.util import _pyyaml_construct_resources, _pyyaml_represent_resources
 from rest.error import HTTPReturn
 from rest.resource import Resource
 from rest.entity.parse import Parser
 from rest.entity.format import Formatter
 
+_pyyaml_construct_resources()
+_pyyaml_represent_resources()
+
 
 class YAMLParser(Parser):
     """Parse an entity in YAML format to native representation."""
 
-    def _construct(self, loader):
-        """Construct a resource node."""
-        def construct(self, node):
-            if isinstance(node, yaml.MappingNode):
-                mapping = loader.construct_mapping(node)
-                resource = Resource(node.tag[1:], mapping)
-            return resource
-        return construct
-
     def parse(self, input, encoding=None):
         """Parse a YAML entity."""
         # We can ignore the encoding as the YAML spec mandates either UTF-8
         # or UTF-16 with a BOM, which can be autodetected.
         # We use a Loader that turns unrecognized !tags into Resources.
-        loader = yaml.Loader(input)
-        loader.add_constructor(None, self._construct(loader))
         try:
-            parsed = loader.get_single_data()
+            parsed = yaml.load(input)
         except YAMLError, e:
             raise HTTPReturn(http.BAD_REQUEST,
-                             reason='YAML Error: %s' % str(e))
+                             reason='YAML load error: %s' % str(e))
         return parsed
 
 
 class YAMLFormatter(Formatter):
     """Format an entity in native representation to YAML."""
 
-    def _represent_resource(self, loader):
-        """Used as a YAML representer for dictionaries."""
-        def format(self, data):
-            if '!type' in data:
-                data = data.copy()
-                tag = '!%s' % data.pop('!type')
-            else:
-                tag = u'tag:yaml.org,2002:map'
-            return loader.represent_mapping(tag, data)
-        return format
-
     def format(self, object, encoding=None):
         """Format a resource as YAML under the specified encoding."""
-        stream = StringIO()
-        dumper = yaml.Dumper(stream, default_flow_style=False, version=(1,1),
-                             encoding=encoding)
-        dumper.add_representer(dict, self._represent_resource(dumper))
-        dumper.add_representer(Resource, self._represent_resource(dumper))
-        dumper.open()
         try:
-            dumper.represent(object)
+            output = yaml.dump(object, default_flow_style=False,
+                               version=(1, 1), encoding=encoding)
         except YAMLError, e:
             raise HTTPReturn(http.INTERNAL_SERVER_ERROR,
                              reason='YAML dump error: %s' % str(e))
-        dumper.close()
-        return stream.getvalue()
+        return output
 
 import sys
 import logging
+import yaml
+
 from rest.api import request
+from rest.resource import Resource
 
 
 def make_absolute(relurl):
     except ImportError:
         return
     return sys.modules[mname]
+
+
+def __construct_resource(loader, node):
+    """YAML constructor that constructs a Resource."""
+    if isinstance(node, yaml.MappingNode):
+        mapping = loader.construct_mapping(node)
+        resource = Resource(node.tag[1:], mapping)
+    else:
+        resource = loader.construct_undefined(node)
+    return resource
+
+def _pyyaml_construct_resources():
+    """Configure PyYAML so that it will process unknown tags and return a
+    Resource instance for them."""
+    yaml.Loader.add_constructor(None, __construct_resource)
+
+
+def __represent_resource(loader, data):
+    if '!type' in data:
+        data = data.copy()
+        tag = '!%s' % data.pop('!type')
+    else:
+        tag = u'tag:yaml.org,2002:map'
+    return loader.represent_mapping(tag, data)
+
+def _pyyaml_represent_resources():
+    """Configure PyYAML such that it will represent a Resource instance with a
+    !tag corresponding to its type."""
+    yaml.Dumper.add_representer(dict, __represent_resource)
+    yaml.Dumper.add_representer(Resource, __represent_resource)