VFSImporter

direct/src/showbase/VFSImporter.py

@@ -0,0 +1,263 @@
+from direct.stdpy.file import open
+from pandac.PandaModules import Filename, VirtualFileSystem, VirtualFileMountSystem
+import sys
+import new
+import os
+import marshal
+import imp
+import struct
+import __builtin__
+
+__all__ = ['register']
+
+vfs = VirtualFileSystem.getGlobalPtr()
+
+# Possible file types.
+FTPythonSource = 0
+FTPythonCompiled = 1
+FTCompiledModule = 2
+
+pycExtension = 'pyc'
+if not __debug__:
+    # In optimized mode, we actually operate on .pyo files, not .pyc
+    # files.
+    pycExtension = 'pyo'
+
+class VFSImporter:
+    """ This class serves as a Python importer to support loading
+    Python .py and .pyc/.pyo files from Panda's Virtual File System,
+    which allows loading Python source files from mounted .mf files
+    (among other places). """
+
+    def __init__(self, path):
+        self.dir_path = Filename.fromOsSpecific(path)
+
+    def find_module(self, fullname):
+        basename = fullname.split('.')[-1]
+        path = Filename(self.dir_path, Filename(basename))
+
+        # First, look for Python files.
+        filename = path
+        filename.setExtension('py')
+        vfile = vfs.getFile(filename, True)
+        if vfile:
+            return VFSLoader(self, vfile, filename, FTPythonSource)
+
+        # If there's no .py file, but there's a .pyc file, load that
+        # anyway.
+        filename = path
+        filename.setExtension(pycExtension)
+        vfile = vfs.getFile(filename, True)
+        if vfile:
+            return VFSLoader(self, vfile, filename, FTPythonCompiled)
+
+        # Look for a compiled C/C++ module.
+        for desc in imp.get_suffixes():
+            if desc[2] != imp.C_EXTENSION:
+                continue
+            
+            filename = path
+            filename.setExtension(desc[0][1:])
+            vfile = vfs.getFile(filename, True)
+            if vfile:
+                return VFSLoader(self, vfile, filename, FTCompiledModule,
+                                 desc = desc)
+        
+
+        # Finally, consider a package, i.e. a directory containing
+        # __init__.py.
+        filename = Filename(path, Filename('__init__.py'))
+        vfile = vfs.getFile(filename, True)
+        if vfile:
+            return VFSLoader(self, vfile, filename, FTPythonSource,
+                             package = True)
+        filename = Filename(path, Filename('__init__.' + pycExtension))
+        vfile = vfs.getFile(filename, True)
+        if vfile:
+            return VFSLoader(self, vfile, filename, FTPythonCompiled,
+                             package = True)
+
+        return None
+
+class VFSLoader:
+    """ The second part of VFSImporter, this is created for a
+    particular .py file or directory. """
+    
+    def __init__(self, importer, vfile, filename, fileType,
+                 desc = None, package = False):
+        self.importer = importer
+        self.dir_path = importer.dir_path
+        self.timestamp = vfile.getTimestamp()
+        self.filename = filename
+        self.fileType = fileType
+        self.desc = desc
+        self.package = package
+    
+    def load_module(self, fullname):
+        if self.fileType == FTCompiledModule:
+            return self._import_compiled_module(fullname)
+        
+        code = self._read_code()
+        if not code:
+            raise ImportError
+        
+        mod = sys.modules.setdefault(fullname, new.module(fullname))
+        mod.__file__ = self.filename.cStr()
+        mod.__loader__ = self
+        if self.package:
+            mod.__path__ = []
+        exec code in mod.__dict__
+        return mod
+
+    def getdata(self, path):
+        path = Filename(self.dir_path, Filename.fromOsSpecific(path))
+        f = open(path, 'rb')
+        return f.read()
+
+    def is_package(self, fullname):
+        return self.package
+
+    def get_code(self, fullname):
+        return self._read_code()
+
+    def get_source(self, fullname):
+        return self._read_source()
+        
+    def _read_source(self):
+        """ Returns the Python source for this file, if it is
+        available, or None if it is not. """
+        
+        if self.fileType == FTPythonCompiled or \
+           self.fileType == FTCompiledModule:
+            return None
+        
+        filename = Filename(self.filename)
+        filename.setExtension('py')
+        try:
+            file = open(filename, 'rU')
+        except IOError:
+            return None
+        return file.read()
+
+    def _import_compiled_module(self, fullname):
+        """ Loads the compiled C/C++ shared object as a Python module,
+        and returns it. """
+
+        print "importing %s" % (fullname)
+
+        vfile = vfs.getFile(self.filename, False)
+
+        # We can only import a compiled module if it already exists on
+        # disk.  This means if it's a truly virtual file that has no
+        # on-disk equivalent, we have to write it to a temporary file
+        # first.
+        if hasattr(vfile, 'getMount') and \
+           isinstance(vfile.getMount(), VirtualFileMountSystem):
+            # It's a real file.
+            filename = self.filename
+        else:
+            # It's a virtual file.  Dump it.
+            filename = Filename.temporary('', self.filename.getBasenameWoExtension(),
+                                          '.' + self.filename.getExtension())
+            filename.setExtension(self.filename.getExtension())
+            fin = open(vfile, 'rb')
+            fout = open(filename, 'wb')
+            data = fin.read(4096)
+            while data:
+                fout.write(data)
+                data = fin.read(4096)
+            fin.close()
+            fout.close()
+
+        module = imp.load_module(fullname, None, filename.toOsSpecific(),
+                                 self.desc)
+        module.__file__ = self.filename.cStr()
+        return module
+        
+
+    def _read_code(self):
+        """ Returns the Python compiled code object for this file, if
+        it is available, or None if it is not. """
+
+        if self.fileType == FTPythonCompiled:
+            # It's a pyc file; just read it directly.
+            pycVfile = vfs.getFile(self.filename, False)
+            if pycVfile:
+                return self._loadPyc(pycVfile)
+            return None
+
+        elif self.fileType == FTCompiledModule:
+            return None
+
+        # It's a .py file (or an __init__.py file; same thing).  Read
+        # the .pyc file if it is available and current; otherwise read
+        # the .py file and compile it.
+        pycFilename = Filename(self.filename)
+        pycFilename.setExtension(pycExtension)
+        pycVfile = vfs.getFile(pycFilename, False)
+        t_pyc = None
+        if pycVfile:
+            t_pyc = pycVfile.getTimestamp()
+
+        code = None
+        if t_pyc and t_pyc >= self.timestamp:
+            code = self._loadPyc(pycVfile)
+
+        if not code:
+            source = self._read_source()
+            filename = Filename(self.filename)
+            filename.setExtension('py')
+            code = self._compile(filename, source)
+
+        return code
+
+    def _loadPyc(self, vfile):
+        """ Reads and returns the marshal data from a .pyc file. """
+        code = None
+        f = open(vfile, 'rb')
+        if f.read(4) == imp.get_magic():
+            t = struct.unpack('<I', f.read(4))[0]
+            if t == self.timestamp:
+                code = marshal.loads(f.read())
+        f.close()
+        return code
+        
+
+    def _compile(self, filename, source):
+        """ Compiles the Python source code to a code object and
+        attempts to write it to an appropriate .pyc file. """
+        
+        if source and source[-1] != '\n':
+            source = source + '\n'
+        code = __builtin__.compile(source, filename.cStr(), 'exec')
+
+        # try to cache the compiled code
+        pycFilename = Filename(filename)
+        pycFilename.setExtension(pycExtension)
+        try:
+            f = open(pycFilename, 'wb')
+        except IOError:
+            pass
+        else:
+            f.write('\0\0\0\0')
+            f.write(struct.pack('<I', self.timestamp))
+            f.write(marshal.dumps(code))
+            f.flush()
+            f.seek(0, 0)
+            f.write(imp.get_magic())
+            f.close()
+
+        return code
+
+_registered = False
+def register():
+    """ Register the VFSImporter on the path_hooks, if it has not
+    already been registered, so that future Python import statements
+    will vector through here (and therefore will take advantage of
+    Panda's virtual file system). """
+    
+    global _registered
+    if not _registered:
+        _registered = True
+        sys.path_hooks.insert(0, VFSImporter)
+