blender_asset_tracer/__init__.py

# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
#  Copyright (C) 2014-2018 Blender Foundation
#
# ##### END GPL LICENSE BLOCK #####

# <pep8 compliant>

__version__ = '1.15'

### --- BAT used as an addon with following code ---

bl_info = {
    "name": "Blender Asset Tracer",
    "author": "Campbell Barton, Sybren A. Stüvel, Loïc Charrière and Clément Ducarteron",
    "version": (1, 15, 0),
    "blender": (2, 80, 0),
    "location": "File > External Data > BAT",
    "description": "Utility for packing blend files",
    "warning": "",
    "wiki_url": "https://developer.blender.org/project/profile/79/",
    "category": "Import-Export",
}


## Reset root module name
## if folder has an unexpected name (like "blender_asset_tracer-main" from zip-dl)
import sys
if __name__ != "blender_asset_tracer":
    sys.modules["blender_asset_tracer"] = sys.modules[__name__]

import bpy
from bpy.types import Operator
from bpy_extras.io_utils import ExportHelper
import zipfile
from blender_asset_tracer.pack import zipped
from pathlib import Path, PurePath

import os
import re
# import sys
import subprocess
import tempfile
from blender_asset_tracer.trace import deps
from . import preferences

class ExportBatPack(Operator, ExportHelper):
    bl_idname = "export_bat.pack"
    bl_label = "Export to Archive using BAT"

    # ExportHelper
    filename_ext = ".zip"

    @classmethod
    def poll(cls, context):
        return bpy.data.is_saved


    def execute(self, context):
        import os
        outfname = bpy.path.ensure_ext(self.filepath, ".zip")
        scn = bpy.context.scene
        
        print(f'outfname {outfname}')
        self.report({'INFO'},'Executing ZipPacker ...')


        with zipped.ZipPacker(
                Path(bpy.data.filepath),
                Path(bpy.data.filepath).parent,
                str(self.filepath)) as packer:
            packer.strategise()
            packer.execute()
        self.report({'INFO'},'Packing successful !')

        with zipfile.ZipFile(str(self.filepath)) as inzip:
            inzip.testzip()

        self.report({'INFO'}, 'Written to %s' % outfname)

        return {'FINISHED'}


class BAT_OT_export_zip(Operator, ExportHelper):
    """Export current blendfile as .ZIP"""
    bl_label = "Zip Pack - keep hierachy" # Export File to .ZIP
    bl_idname = "bat.export_zip"
    
    filename_ext = '.zip'
    
    
    root_dir : bpy.props.StringProperty(
        name="Root",
        description='Top Level Folder of your project.'
        '\nFor now Copy/Paste correct folder by hand if default is incorrect.'
        '\n!!! Everything outside won\'t be zipped !!!',
    )

    @classmethod
    def poll(cls, context):
        return bpy.data.is_saved

    ## ExportHelper has it's own invoke so overriding it with this one create error
    # def invoke(self, context, event):
    #     prefs = preferences.get_addon_prefs()
    #     if prefs.root_default:
    #         self.root_dir = prefs.root_default
    #     return self.execute(context)

    def execute(self, context):

        root_dir = self.root_dir
        print('root_dir: ', root_dir)

        def open_folder(folderpath):
            """
            open the folder at the path given
            with cmd relative to user's OS
            """
            from shutil import which
            my_os = sys.platform
            if my_os.startswith(('linux','freebsd')):
                cmd = 'xdg-open'
            elif my_os.startswith('win'):
                cmd = 'explorer'
                if not folderpath:
                    return('/')
            else:
                cmd = 'open'

            if not folderpath:
                return('//')

            if os.path.isfile(folderpath): # When pointing to a file
                select = False
                if my_os.startswith('win'):
                    # Keep same path but add "/select" the file (windows cmd option)
                    cmd = 'explorer /select,'
                    select = True

                elif my_os.startswith(('linux','freebsd')):
                    if which('nemo'):
                        cmd = 'nemo --no-desktop'
                        select = True
                    elif which('nautilus'):
                        cmd = 'nautilus --no-desktop'
                        select = True

                if not select:
                    # Use directory of the file
                    folderpath = os.path.dirname(folderpath)

            folderpath = os.path.normpath(folderpath)
            fullcmd = cmd.split() + [folderpath]
            # print('use opening command :', fullcmd)
            subprocess.Popen(fullcmd)        


        def zip_with_structure(zip, filelist, root=None, compressed=True):
            '''
            Zip passed filelist into a zip with root path as toplevel tree
            If root is not passed, the shortest path in filelist becomes the root

            :zip: output fullpath of the created zip
            :filelist: list of filepaht as string or Path object (converted anyway)
            :root: top level of the created hierarchy (not included), file that are not inside root are discarded
            :compressed: Decide if zip is compressed or not
            '''

            filelist = [Path(f) for f in filelist] # ensure pathlib
            if not filelist:
                return
            if not root:
                # autodetect the path thats is closest to root
                #root = sorted(filelist, key=lambda f: f.as_posix().count('/'))[0].parent
                filelist_abs = [str(fl) for fl in filelist]
                root = Path(os.path.commonpath(filelist_abs))
                #print('root: ', root)
            else:
                root = Path(root)

            compress_type = zipfile.ZIP_DEFLATED if compressed else zipfile.ZIP_STORED
            with zipfile.ZipFile(zip, 'w',compress_type) as zipObj:        
                for f in filelist:
                    #print('f: ', f, type(f))
                    if not f.exists():
                        print(f'Not exists: {f.name}')
                        continue
                    if str(root) not in str(f):
                        print(f'{f} is out of root {root}')
                        continue

                    ## 
                    arcname = f.as_posix().replace(root.as_posix(), '').lstrip('/')
                    print(f'adding: {arcname}')
                    zipObj.write(f, arcname)

            return zip, root


        links = []

        current_file = Path(bpy.data.filepath)
        links.append(str(current_file))
        file_link = list(deps(current_file))
        for l in file_link:
            if Path(l.abspath).exists() == False:
                continue

            links.append(l.abspath)
            if l.is_sequence:            
                split_path = PurePath(l.abspath).parts
                file_dir = os.path.join(*split_path[:-1])
                file_name = split_path[-1]

                pattern = '[0-9]+\.[a-zA-Z]+$'
                file_name = re.sub(pattern, '', file_name)
                
                for im in os.listdir(Path(f"{file_dir}/")):
                    if im.startswith(file_name) and re.search(pattern, im):
                        links.append(os.path.join(file_dir, im))

        links = list(set(links))
  
        #output_name = current_file.name
        output = self.filepath
        #print('output: ', output)
        
        root_dir = zip_with_structure(output, links, root_dir)
        root_dir = str(root_dir[1])

        log_output = Path(tempfile.gettempdir(),'README.txt')
        with open(log_output, 'w') as log:
            log.write("File is located here:")
            log.write(f"\n   - /{str(current_file).replace(root_dir,'')}")
        
        with zipfile.ZipFile(output, 'a') as zipObj:
            zipObj.write(log_output, log_output.name)
        
        ### Same as
        #zipObj = zipfile.ZipFile(output, 'a')
        #zipObj.write(log_output, log_output.name)
        #zipObj.close()

        open_folder(Path(output).parent)

        return {'FINISHED'}


def menu_func(self, context):
    layout = self.layout
    layout.separator()
    layout.operator(ExportBatPack.bl_idname)
    filepath = layout.operator(BAT_OT_export_zip.bl_idname)

    prefs = preferences.get_addon_prefs()
    
    root_dir_env = None
    if prefs.use_env_root:
        root_dir_env = os.getenv('ZIP_ROOT')
        if not root_dir_env: # first fallback to PROJECT_ROOT
            root_dir_env = os.getenv('PROJECT_ROOT')
        
        if not root_dir_env: # if no env, use prefs instead
            root_dir_env = prefs.root_default
    
    filepath.root_dir = '' if root_dir_env == None else root_dir_env


classes = (
    ExportBatPack,
    BAT_OT_export_zip,
)


def register():
    preferences.register()
    for cls in classes:
        bpy.utils.register_class(cls)

    bpy.types.TOPBAR_MT_file_external_data.append(menu_func)


def unregister():
    for cls in classes:
        bpy.utils.unregister_class(cls)

    preferences.unregister()
    bpy.types.TOPBAR_MT_file_external_data.remove(menu_func)


if __name__ == "__main__":
    register()