Commits

Anonymous committed 8229e14

cleanup

Comments (0)

Files changed (1)

nbt2yaml/yamlgen.py

 import yaml
 from nbt2yaml import parse
 
-explicit_types = parse.TAG_Short, parse.TAG_Long, parse.TAG_Double, parse.TAG_Byte, parse.TAG_Byte_Array
+_explicit_types = parse.TAG_Short, parse.TAG_Long, parse.TAG_Double, parse.TAG_Byte, parse.TAG_Byte_Array
+
+_canned_types = {
+    str:parse.TAG_String,
+    float:parse.TAG_Float,
+    int:parse.TAG_Int,
+}
+
+_types_to_python = {
+    parse.TAG_Long:long,
+    parse.TAG_Int:int,
+    parse.TAG_Short:int,
+    parse.TAG_Byte:int,
+    parse.TAG_Float:float,
+    parse.TAG_Double:float
+}
+
 class ForceType(object):
     """Represent a data value with an explicit type.
     
         return ForceType(type_, value)
     return _constructor
 
-for type_ in explicit_types:
+for type_ in _explicit_types:
     yaml.add_constructor(u'!%s' % type_.name, _type_constructor(type_))
 
-def yaml_serialize(struct):
+def _yaml_serialize(struct):
     tag, name, data = struct.type, struct.name, struct.data
     name = name.encode('utf-8')
     return {name:_value_as_yaml(tag, data)}
 
+def _yaml_deserialize(struct):
+    name, data = struct.items()[0]
+    type_ = _type_from_yaml(data)
+    return parse.Tag._tuple(type_, name.decode('utf-8'), _yaml_as_value(type_, data))
+
 def _value_as_yaml(type_, value):
     if type_ is parse.TAG_Compound:
-        return [yaml_serialize(s) for s in value]
+        return [_yaml_serialize(s) for s in value]
     elif type_ is parse.TAG_String:
         return value.encode('utf-8')
-    elif type_ in explicit_types:
+    elif type_ in _explicit_types:
         return ForceType(type_, value)
     elif type_ is parse.TAG_List:
         element_type, data = value
 
 def _yaml_as_value(type_, value):
     if type_ is parse.TAG_Compound:
-        return [yaml_deserialize(s) for s in value]
+        return [_yaml_deserialize(s) for s in value]
     elif type_ is parse.TAG_String:
         return value.decode('utf-8')
-    elif type_ in explicit_types:
-        if type_ in (parse.TAG_Long,):
-            return long(value.value)
-        elif type_ in (parse.TAG_Int, parse.TAG_Short, parse.TAG_Byte):
-            return int(value.value)
-        elif type_ in (parse.TAG_Float, parse.TAG_Double):
-            return float(value.value)
+    elif type_ in _explicit_types:
+        if type_ in _types_to_python:
+            return _types_to_python[type_](value.value)
         else:
             return value.value
     elif type_ is parse.TAG_List:
 def _type_from_yaml(data):
     if isinstance(data, list):
         if isinstance(data[0], dict):
-            type_ = parse.TAG_Compound
+            return parse.TAG_Compound
         else:
-            type_ = parse.TAG_List
+            return parse.TAG_List
     elif isinstance(data, ForceType):
-        type_ = data.type
-    elif type(data) in canned_types:
-        type_ = canned_types[type(data)]
+        return data.type
+    elif type(data) in _canned_types:
+        return _canned_types[type(data)]
     else:
-        raise ValueError("Can't determine type for element: %r" % (data))
-    return type_
-
-canned_types = {
-    str:parse.TAG_String,
-    float:parse.TAG_Float,
-    int:parse.TAG_Int,
-}
-
-def yaml_deserialize(struct):
-    name, data = struct.items()[0]
-    type_ = _type_from_yaml(data)
-
-    return parse.Tag._tuple(type_, name.decode('utf-8'), _yaml_as_value(type_, data))
+        raise ValueError("Can't determine type for element: %r" % data)
 
 def dump_yaml(struct, canonical=False, default_flow_style=False):
+    """Dump a yaml string given an nbt structure parsed by nbt2yaml."""
+
     return yaml.dump(
-                yaml_serialize(struct), 
+                _yaml_serialize(struct), 
                 default_flow_style=default_flow_style, 
                 canonical=canonical)
 
 
 def parse_yaml(stream):
-    struct = yaml.load(stream)
-    return yaml_deserialize(struct)
+    """Parse a yaml stream into an nbt structure readable by nbt2yaml."""
+
+    return _yaml_deserialize(yaml.load(stream))