Merge pull request #3 from swillisart/python3compat

Python 2/3 compatibility

maya_scenefile_parser/binary.py

@@ -1,48 +1,55 @@
 import os
+import sys
 import struct
 
 from . import common, iff
 
+if sys.version_info < (3, 0):
+    def makeBytes(data):
+        return bytes(data)
+else:
+    def makeBytes(data):
+        return bytes(data, 'utf8')
 
 # IFF chunk type IDs
-FOR4 = common.be_word4("FOR4")
-LIS4 = common.be_word4("LIS4")
+FOR4 = common.be_word4(b"FOR4")
+LIS4 = common.be_word4(b"LIS4")
 # 64 bits
-FOR8 = common.be_word4("FOR8")
-LIS8 = common.be_word4("LIS8")
+FOR8 = common.be_word4(b"FOR8")
+LIS8 = common.be_word4(b"LIS8")
 
 # General
-MAYA = common.be_word4("Maya")
+MAYA = common.be_word4(b"Maya")
 
 # File referencing
-FREF = common.be_word4("FREF")
-FRDI = common.be_word4("FRDI")
+FREF = common.be_word4(b"FREF")
+FRDI = common.be_word4(b"FRDI")
 
 # Header fields
-HEAD = common.be_word4("HEAD")
-VERS = common.be_word4("VERS")
-PLUG = common.be_word4("PLUG")
-FINF = common.be_word4("FINF")
-AUNI = common.be_word4("AUNI")
-LUNI = common.be_word4("LUNI")
-TUNI = common.be_word4("TUNI")
+HEAD = common.be_word4(b"HEAD")
+VERS = common.be_word4(b"VERS")
+PLUG = common.be_word4(b"PLUG")
+FINF = common.be_word4(b"FINF")
+AUNI = common.be_word4(b"AUNI")
+LUNI = common.be_word4(b"LUNI")
+TUNI = common.be_word4(b"TUNI")
 
 # Node creation
-CREA = common.be_word4("CREA")
-SLCT = common.be_word4("SLCT")
-ATTR = common.be_word4("ATTR")
+CREA = common.be_word4(b"CREA")
+SLCT = common.be_word4(b"SLCT")
+ATTR = common.be_word4(b"ATTR")
 
-CONS = common.be_word4("CONS")
-CONN = common.be_word4("CONN")
+CONS = common.be_word4(b"CONS")
+CONN = common.be_word4(b"CONN")
 
 # Data types
-FLGS = common.be_word4("FLGS")
-DBLE = common.be_word4("DBLE")
-DBL3 = common.be_word4("DBL3")
-STR_ = common.be_word4("STR ")
-FLT2 = common.be_word4("FLT2")
-CMPD = common.be_word4("CMPD")
-MESH = common.be_word4("MESH")
+FLGS = common.be_word4(b"FLGS")
+DBLE = common.be_word4(b"DBLE")
+DBL3 = common.be_word4(b"DBL3")
+STR_ = common.be_word4(b"STR ")
+FLT2 = common.be_word4(b"FLT2")
+CMPD = common.be_word4(b"CMPD")
+MESH = common.be_word4(b"MESH")
 
 
 MAYA_BINARY_32 = iff.IffFormat(
@@ -66,14 +73,15 @@ class MayaBinaryError(RuntimeError):
 
 
 class MayaBinaryParser(iff.IffParser, common.MayaParserBase):
-    def __init__(self, stream):
+    def __init__(self, stream, info_only=False):
         # Determine Maya format based on magic number
         # Maya 2014+ files begin with a FOR8 block, indicating a 64-bit format.
+        self.info_only = info_only
         magic_number = stream.read(4)
         stream.seek(0)
-        if magic_number == "FOR4":
+        if magic_number == b"FOR4":
             format = MAYA_BINARY_32
-        elif magic_number == "FOR8":
+        elif magic_number == b"FOR8":
             format = MAYA_BINARY_64
         else:
             raise MayaBinaryError("Bad magic number")
@@ -91,21 +99,24 @@ def __init__(self, stream):
         self._load_mtypeid_database(os.path.join(os.path.dirname(__file__),
                                     "modules", "maya", "2012", "typeids.dat"))
 
+
     def on_iff_chunk(self, chunk):
         if chunk.typeid == self.__node_chunk_type:
             mtypeid = self._read_mtypeid()
             if mtypeid == MAYA:
                 self._handle_all_chunks()
             elif mtypeid == HEAD:
                 self._parse_maya_header()
+            elif self.info_only:
+                return
             elif mtypeid == FREF:
                 self._parse_file_reference()
             elif mtypeid == CONN:
                 self._parse_connection()
             else:
                 self._parse_node(mtypeid)
 
-        elif chunk.typeid == self.__list_chunk_type:
+        elif chunk.typeid == self.__list_chunk_type and not self.info_only:
             mtypeid = self._read_mtypeid()
             if mtypeid == CONS:
                 self._handle_all_chunks()
@@ -183,7 +194,7 @@ def _parse_node(self, mtypeid):
             # Create node
             if chunk.typeid == CREA:
                 typename = self.__mtypeid_to_typename.get(mtypeid, "unknown")
-                name_parts = self._read_chunk_data(chunk)[1:-1].split("\0")
+                name_parts = self._read_chunk_data(chunk)[1:-1].split(b"\0")
                 name = name_parts[0]
                 parent_name = name_parts[1] if len(name_parts) > 1 else None
                 self.on_create_node(typename, name, parent=parent_name)
@@ -247,7 +258,7 @@ def _load_mtypeid_database(self, path):
         with open(path) as f:
             line = f.readline()
             while line:
-                mtypeid = common.be_word4(line[:4])
+                mtypeid = common.be_word4(makeBytes(line[:4]))
                 typename = line[5:].strip()
                 self.__mtypeid_to_typename[mtypeid] = typename
                 line = f.readline()

maya_scenefile_parser/common.py

@@ -2,57 +2,57 @@
 
 
 def be_word4(buf):
-    return struct.unpack(">L", buf)[0]
+    return struct.unpack('>L', buf)[0]
 
 
 def le_word4(buf):
-    return struct.unpack("<L", buf)[0]
+    return struct.unpack('<L', buf)[0]
 
 
 def be_word8(buf):
-    return struct.unpack(">Q", buf)[0]
+    return struct.unpack('>Q', buf)[0]
 
 
 def le_word8(buf):
-    return struct.unpack("<Q", buf)[0]
+    return struct.unpack('<Q', buf)[0]
 
 
 def be_read4(stream):
-    return struct.unpack(">L", stream.read(4))[0]
+    return struct.unpack('>L', stream.read(4))[0]
 
 
 def le_read4(stream):
-    return struct.unpack("<L", stream.read(4))[0]
+    return struct.unpack('<L', stream.read(4))[0]
 
 
 def be_read8(stream):
-    return struct.unpack(">Q", stream.read(8))[0]
+    return struct.unpack('>Q', stream.read(8))[0]
 
 
 def le_read8(stream):
-    return struct.unpack("<Q", stream.read(8))[0]
+    return struct.unpack('<Q', stream.read(8))[0]
 
 
 def align(size, stride):
     return stride * int(1 + ((size - 1) / stride))
 
 
 def read_null_terminated(stream):
-    result = ""
+    result = b''
     next = stream.read(1)
-    while stream and next != '\0':
+    while stream and next != b'\0':
         result += next
         next = stream.read(1)
     return result
 
 
 def plug_element_count(plug):
-    lbracket = plug.rfind("[")
+    lbracket = plug.rfind(b'[')
     if lbracket != -1:
-        rbracket = plug.rfind("]")
+        rbracket = plug.rfind(b']')
         if rbracket != -1 and lbracket < rbracket:
             slicestr = plug[lbracket + 1:rbracket]
-            bounds = slicestr.split(":")
+            bounds = slicestr.split(b':')
             if len(bounds) > 1:
                 return int(bounds[1]) - int(bounds[0]) + 1
     return 1