vse_toolbox/operators/imports.py

from pathlib import Path

import bpy
from bpy.types import Operator, UIList
from bpy.props import (CollectionProperty, BoolProperty, EnumProperty, StringProperty)

from vse_toolbox.constants import (EDITS, EDIT_SUFFIXES, MOVIES, MOVIE_SUFFIXES,
    SOUNDS, SOUND_SUFFIXES)

from vse_toolbox.sequencer_utils import (clean_sequencer, import_edit, import_movie,
    import_sound, get_strips, get_channel_index)
       
from vse_toolbox.bl_utils import get_scene_settings
from vse_toolbox.file_utils import install_module


class VSETB_OT_auto_select_files(Operator):
    bl_idname = "vse_toolbox.auto_select_files"
    bl_label = "Auto Select"
    bl_description = "Auto Select Files"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return True

    def get_items(self, items=[]):
        if not items:
            return [('NONE', 'None', '', 0)]
        return [(e, e, '', i) for i, e in enumerate(items)]

    def execute(self, context):
        params = context.space_data.params
        directory = Path(params.directory.decode())
        
        EDITS.clear()
        MOVIES.clear()
        SOUNDS.clear()
        
        edits = []
        movies = []
        sounds = []

        for file_entry in directory.glob('*'):
            if file_entry.is_dir():
                continue

            if file_entry.suffix in EDIT_SUFFIXES:
                edits.append(file_entry.name)
            elif file_entry.suffix in MOVIE_SUFFIXES:
                movies.append(file_entry.name)
            elif file_entry.suffix in SOUND_SUFFIXES:
                sounds.append(file_entry.name)

        edits.sort(reverse=True)
        movies.sort(reverse=True)
        sounds.sort(reverse=True)

        EDITS.extend(self.get_items(items=edits))
        MOVIES.extend(self.get_items(items=movies))
        SOUNDS.extend(self.get_items(items=sounds))

        return {'FINISHED'}


class VSETB_OT_import_files(Operator):
    bl_idname = "vse_toolbox.import_files"
    bl_label = "Import"
    bl_description = "Import Edit"
    bl_options = {"REGISTER", "UNDO"}
    
    directory : StringProperty(subtype='DIR_PATH')
    filepath: StringProperty(
        name="File Path",
        description="Filepath used for importing the file",
        maxlen=1024,
        subtype='FILE_PATH',
    )
    files : CollectionProperty(type=bpy.types.OperatorFileListElement)
    clean_sequencer : BoolProperty(
        name="Clean Sequencer",
        default=False,
        description="Clean all existing strips in sequencer",
    )
    
    import_edit : BoolProperty(name='', default=True)
    edit: EnumProperty(name='', items=lambda s, c: EDITS)
    match_by : EnumProperty(name='Match By', items=[('NAME', 'Name', ''), ('INDEX', 'Index', '')])

    import_movie : BoolProperty(name='', default=False)
    movie: EnumProperty(name='', items=lambda s, c: MOVIES)

    import_sound : BoolProperty(name='', default=False)
    sound: EnumProperty(name='', items=lambda s, c: SOUNDS)

    @classmethod
    def poll(cls, context):
        return True

    def draw(self, context):
        scn = context.scene
        settings = get_scene_settings()

        layout = self.layout
        layout.use_property_split = True 
        layout.use_property_decorate = False

        col = layout.column(align=True)
        col.operator('vse_toolbox.auto_select_files', text='Auto Select')

        row = layout.row(heading="Import Edit", align=True)
        row.prop(self, 'import_edit')
        sub = row.row(align=True)
        sub.active = self.import_edit
        sub.prop(self, 'edit')
        row = layout.row(align=True)
        row.prop(self, 'match_by', expand=True)

        layout.separator()

        row = layout.row(heading="Import Movie", align=True)
        row.prop(self, 'import_movie')
        sub = row.row()
        sub.active = self.import_movie
        sub.prop(self, 'movie')

        row = layout.row(heading="Import Sound", align=True)
        row.prop(self, 'import_sound')
        sub = row.row()
        sub.active = self.import_sound
        sub.prop(self, 'sound')

        col = layout.column()
        col.separator()
        col.prop(self, 'clean_sequencer')

    def invoke(self, context, event):
        context.window_manager.fileselect_add(self)
        return {'RUNNING_MODAL'}

    def execute(self, context):
        sequencer = context.scene.sequence_editor.sequences

        edit_filepath = Path(self.directory, self.edit)
        if not edit_filepath.exists():
            self.import_edit = False

        movie_filepath = Path(self.directory, self.movie)
        if not movie_filepath.exists():
            self.import_movie = False
        
        sound_filepath = Path(self.directory, self.sound)
        if not sound_filepath.exists():
            self.import_sound = False
        
        if self.clean_sequencer:
            clean_sequencer(
                edit=self.import_edit,
                movie=self.import_movie,
                sound=self.import_sound,
            )
        
        if self.import_edit:
            print(f'[>.] Loading Edit from: {str(edit_filepath)}')

            import_edit(edit_filepath, adapter="cmx_3600", match_by=self.match_by)
        
        if self.import_movie:
            print(f'[>.] Loading Movie from: {str(movie_filepath)}')

            for strip in get_strips(channel='Movie'):
                sequencer.remove(strip)

            import_movie(movie_filepath)

        if self.import_sound or (not self.import_sound and self.import_movie):

            for strip in get_strips(channel='Audio'):
                sequencer.remove(strip)

            if self.import_sound:
                print(f'[>.] Loading Audio from: {str(sound_filepath)}')
                import_sound(sound_filepath)
            else:
                print(f'[>.] Loading Audio from: {str(movie_filepath)}')
                import_sound(movie_filepath)
        
        context.scene.sequence_editor.sequences.update()

        return {"FINISHED"}

class VSETB_UL_import_task(UIList):
    def draw_item(self, context, layout, data, item, icon, active_data,
                  active_propname, index):

        layout.prop(item, 'import_enabled', text='')
        layout.label(text=item.name)
        
class VSETB_OT_import_shots(Operator):
    bl_idname = "vse_toolbox.import_shots"
    bl_label = "Import Shots"
    bl_description = "Import Shots for disk or tracker"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return True

    def execute(self, context):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project

        from gadget import Project

        p = Project.from_env()
        p.shots.fetch()
        
        tasks = [t for t in project.task_types if t.import_enabled]
        for i, task_type in enumerate(tasks):
            channel_index = 9 + i

            scn.sequence_editor.channels[channel_index].name = task_type.name

        # Create Channels:
        scn.sequence_editor.channels[9]

        for strip in get_strips('Shots'):
            for i, task_type in enumerate(tasks):
                shot = p.shots[strip.name]
                versions = shot.tasks[task_type.name].outputs[0].versions.fetch()

                if versions:
                    print(versions[-1])
                
                    imported_strip = import_movie(versions[-1].path)
                    imported_strip.frame_start = strip.frame_start

                    if imported_strip.frame_final_duration != strip.frame_final_duration:
                        print(f'strip {strip.name} has different duration')
                        imported_strip.color_tag = 'COLOR_03'
                    
                    imported_strip.frame_final_end = strip.frame_final_end

                    imported_strip.channel = get_channel_index(task_type.name)
                    print(get_channel_index(task_type.name), imported_strip.channel)
                else:
                    print(f'No movie for {strip.name} of task {task_type.name}')

        # for task_type in tasks:
        #     for strip in get_strips(task_type.name):
        #         strip.channel = get_channel_index(task_type.name)


        return {'FINISHED'}

    def draw(self, context):
        settings = get_scene_settings()
        project = settings.active_project

        layout = self.layout
        col = layout.column(align=False)
        col.use_property_split = True
        col.use_property_decorate = False

        row = col.row(align=True)
        row.prop(project, "import_source", text='Source', expand=True)

        row = col.row(align=True)
        row.prop(project, 'import_task', text='Import Tasks')

        if project.import_task == 'SELECTED':
            col.use_property_split = False
            col.template_list("VSETB_UL_import_task", "import_task", project, "task_types", project, "task_type_index", rows=8)

    def invoke(self, context, event):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project
        
        return context.window_manager.invoke_props_dialog(self)

classes = (
    VSETB_UL_import_task,
    VSETB_OT_auto_select_files,
    VSETB_OT_import_files,
    VSETB_OT_import_shots,
)

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

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