Modified ZipFile() to support writedir() (with auto pyc/py selection)

class _ZipFile ( zipfile.ZipFile ):
    # Inspired by:
    # http://www.calazan.com/how-to-zip-an-entire-directory-with-python/comment-page-1/#comment-45304
    #
    # Heavily modified to support:
    #  - automatic choice between .pyc and .py
    #  - partially integrated with ZipFile class
    #
    # TODO:
    #  - Allow auto compile to /tmp when .pyc is not present.
    
    def writedir(self, pathname, basename = ""):
        """Zip the contents of an entire folder (with that folder included
        in the archive). Empty subfolders will be included in the archive
        as well.
        """
        
        _paths = {}
        
        def debug_log(msg):
            if self.debug:
                print msg
        
        def AddToZIP(absolute_path, relative_path):
            # check the ext of the file
            _path, _ext = os.path.splitext(absolute_path)
            
            # check if it is a file
            _is_file = os.path.isfile(absolute_path)
            
            # generate absolute pyc/py paths
            _pyc = "%s.pyc" % ( _path )
            _py = "%s.py" % ( _path )
            
            # generate relative pyc/py paths
            _pyc_rel = "%s.pyc" % ( _path )
            _py_rel = "%s.py" % ( _path )
            
            if _is_file:
                # if .py, then check if pyc is there
                if _ext == '.py':
                    # if its there then dont add
                    if _paths.has_key(_pyc):
                        debug_log("[py/skip] '%s'" % ( _py ) )
                    else:
                        _paths[_py] = _py_rel
                
                # if .pyc, then check if we have a py
                elif _ext == '.pyc':
                    # if we have a py, then replace it with the pyc
                    if _paths.has_key(_py):
                        debug_log("[pyc/replace] '%s'" % ( _pyc) )
                        _paths[_pyc] = _pyc_rel
                    else:
                        debug_log("[pyc/add] '%s'" % ( _pyc ) )
                        _paths[_pyc] = _pyc_rel
                        
            else:
                debug_log("[dir/add] '%s'" % ( absolute_path ) )
                _paths[absolute_path] = relative_path
                
        def GenerateZIP():
            for _abs, _rel in _paths.iteritems():
                print "[zip/add] %s" % ( _abs )
                self.write(_abs, _rel)
            
        parent_folder = os.path.dirname(pathname)
        # Retrieve the paths of the folder contents.
        contents = os.walk(pathname)
        
        try:            
            for root, folders, files in contents:
                # Include all subfolders, including empty ones.
                for folder_name in folders:
                    absolute_path = os.path.join(root, folder_name)
                    relative_path = absolute_path.replace(parent_folder + '\\','')
                    AddToZIP(absolute_path, relative_path)
                    
                for file_name in files:
                    absolute_path = os.path.join(root, file_name)
                    relative_path = absolute_path.replace(parent_folder + '\\','')
                    AddToZIP(absolute_path, relative_path)
                    
            GenerateZIP()
            
            debug_log("ZIP generated successfully.")
            
        except IOError, message:
            print message
            raise
        
        except OSError, message:
            print message
            raise
        
        except zipfile.BadZipfile, message:
            print message
            raise
        
        finally:
            self.close()
            
# usage
# create zip object
_zip_obj = _ZipFile("filenamehere.txt", 'w')
        
# enable debug mode
_zip_obj.debug = debug
        
# attempt to zip the dir
_zip_obj.writedir(
    pathname = dirpath
)