vse_toolbox/operators/spreadsheet.py

742 lines
27 KiB
Python

# SPDX-License-Identifier: GPL-2.0-or-later
import csv
from datetime import datetime
import os
from pathlib import Path
from io import StringIO
import bpy
from bpy.types import (Operator, Menu, OperatorFileListElement)
from bpy.props import (EnumProperty, StringProperty, CollectionProperty)
from bl_operators.presets import AddPresetBase
from vse_toolbox.sequencer_utils import (get_strips, get_strip_sequence_name, get_channel_index)
from vse_toolbox.bl_utils import (get_addon_prefs, get_scene_settings)
from vse_toolbox.file_utils import (open_file, install_module, fuzzy_match, norm_str)
from vse_toolbox.constants import SPREADSHEET
class VSETB_MT_export_spreadsheet_presets(Menu):
bl_label = 'Presets'
preset_subdir = 'vse_toolbox'
preset_operator = 'script.execute_preset'
draw = Menu.draw_preset
class VSETB_OT_add_export_spreadsheet_preset(AddPresetBase, Operator):
bl_idname = "vse_toolbox.add_spreadsheet_preset"
bl_label = "Add Spreadsheet Preset"
bl_description = "Add Spreadsheet Preset"
bl_options = {"REGISTER", "UNDO"}
preset_menu = 'VSETB_MT_export_spreadsheet_presets'
#preset_menu = 'VSETB_OT_MT_spreadsheet_presets'
# Common variable used for all preset values
#C.scene.vsetb_settings.active_project.spreadsheet_options
preset_defines = [
'scene = bpy.context.scene',
'settings = scene.vsetb_settings',
'project = settings.active_project'
]
# Properties to store in the preset
preset_values = [
'project.spreadsheet_export',
]
# Directory to store the presets
preset_subdir = 'vse_toolbox/spreadsheet_export'
class VSETB_OT_spreadsheet_cell_move(Operator):
bl_idname = "vse_toolbox.spreadsheet_move"
bl_label = "Move Spreadsheet items"
bl_description = "Move Spreadsheet items"
bl_options = {"REGISTER", "UNDO"}
direction: EnumProperty(
items=(
('UP', "Up", ""),
('DOWN', "Down", ""),
)
)
def execute(self, context):
scn = context.scene
project = get_scene_settings().active_project
spreadsheet = project.spreadsheet_export
cells = spreadsheet.cells
idx = spreadsheet.cell_index
try:
item = cells[idx]
except IndexError:
pass
else:
if self.direction == 'DOWN' and idx < len(cells) - 1:
item_next = cells[idx+1].name
cells.move(idx, idx+1)
spreadsheet.cell_index += 1
elif self.direction == 'UP' and idx >= 1:
item_prev = cells[idx-1].name
cells.move(idx, idx-1)
spreadsheet.cell_index -= 1
return {"FINISHED"}
class VSETB_OT_spreadsheet_from_file(Operator):
bl_idname = "vse_toolbox.spreadsheet_from_file"
bl_label = "Read Spreadsheet Column"
bl_description = "Read Spreadsheet Column"
bl_options = {"REGISTER", "UNDO"}
filepath : StringProperty()
def execute(self, context):
scn = context.scene
project = get_scene_settings().active_project
spreadsheet = project.spreadsheet_import
separator = spreadsheet.separator.replace('\\n', '\n').replace('\\t', '\t').replace('\\r', '\r')
delimiter = spreadsheet.delimiter.replace('\\n', '\n').replace('\\t', '\t').replace('\\r', '\r')
import_cells = spreadsheet.cells
import_cells.clear()
if not self.filepath:
self.report({'ERROR'}, 'No filepath provided')
return {'CANCELLED'}
if not Path(self.filepath).exists():
self.report({'ERROR'}, f'Filepath {self.filepath} not exists')
return {'CANCELLED'}
filepath = Path(self.filepath)
if filepath.suffix.lower() == '.csv':
if len(delimiter) == 1:
with open(filepath, newline=separator) as csvfile:
rows = list(csv.reader(csvfile, delimiter=delimiter))
else: # Auto detect
with open(filepath, 'r', newline="") as csvfile:
dialect = csv.Sniffer().sniff(csvfile.read(1024), delimiters=delimiter)
csvfile.seek(0)
rows = list(csv.reader(csvfile, dialect))
elif filepath.suffix.lower() == '.xlsx':
try:
import openpyxl
except ModuleNotFoundError:
self.report({'INFO'}, 'Installing openpyxl')
openpyxl = install_module('openpyxl')
from openpyxl import Workbook
workbook = openpyxl.load_workbook(filepath, read_only=True)
sheet = workbook.active
rows = [[(c.value or '') for c in r] for r in sheet.rows]
workbook.close()
else:
self.report({'ERROR'}, f'File extension {filepath.suffix} should be in [.csv, .xlsx]')
return {'CANCELLED'}
rows = [r for r in rows if any(r)]
cell_types = project.get_cell_types()
for cell_name in rows[0]:
if not cell_name:
continue
cell = import_cells.add()
cell.name = cell_name
if cell_types.get(cell_name):
cell.import_name = cell_name
continue
matches = [(k, fuzzy_match(cell_name, k)) for k in cell_types.keys()]
best_name, best_match = max(matches, key=lambda x: x[1])
cell.import_name = best_name
cell.enabled = best_match > 0.85
project.spreadsheet_import.use_custom_cells = True
SPREADSHEET.extend(rows)
return {"FINISHED"}
class VSETB_OT_spreadsheet_from_clipboard(Operator):
bl_idname = "vse_toolbox.spreadsheet_from_clipboard"
bl_label = "Read Spreadsheet from clipboard"
bl_description = "Read Spreadsheet from clipboard"
bl_options = {"REGISTER", "UNDO"}
def execute(self, context):
scn = context.scene
project = get_scene_settings().active_project
import_cells = project.spreadsheet_import.cells
import_cells.clear()
SPREADSHEET.clear()
spreadsheet = context.window_manager.clipboard
cell_types = project.get_cell_types()
rows = [r for r in csv.reader(StringIO(spreadsheet), delimiter='\t') if any(r)]
for cell_name in rows[0]:
if not cell_name:
continue
cell = import_cells.add()
cell.name = cell_name
cell.import_name = max(cell_types.keys(), key=lambda x: fuzzy_match(cell_name, x))
cell.enabled = True
project.spreadsheet_import.use_custom_cells = True
SPREADSHEET.extend(rows)
return {"FINISHED"}
class VSETB_OT_spreadsheet_to_clipboard(Operator):
bl_idname = "vse_toolbox.spreadsheet_to_clipboard"
bl_label = "Copy Spreadsheet to clipboard"
bl_description = "Copy Spreadsheet to clipboard"
bl_options = {"REGISTER", "UNDO"}
def execute(self, context):
scn = context.scene
project = get_scene_settings().active_project
import_cells = project.spreadsheet_import.cells
import_cells.clear()
SPREADSHEET.clear()
spreadsheet = context.window_manager.clipboard
cell_types = project.get_cell_types()
rows = list(csv.reader(StringIO(spreadsheet), delimiter='\t'))
for cell_name in rows[0]:
if not cell_name:
continue
cell = import_cells.add()
cell.name = cell_name
cell.import_name = max(cell_types.keys(), key=lambda x: fuzzy_match(cell_name, x))
cell.enabled = True
project.spreadsheet_import.use_custom_cells = True
SPREADSHEET.extend(rows)
return {"FINISHED"}
class VSETB_OT_import_spreadsheet(Operator):
bl_idname = "vse_toolbox.import_spreadsheet"
bl_label = "Import Spreadsheet"
bl_description = "Create strips from nb frames with casting and custom data"
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=OperatorFileListElement)
@classmethod
def poll(cls, context):
settings = get_scene_settings()
return settings.active_project
def invoke(self, context, event):
settings = get_scene_settings()
project = settings.active_project
context.window_manager.fileselect_add(self)
return {'RUNNING_MODAL'}
def draw(self, context):
scn = context.scene
settings = get_scene_settings()
project = settings.active_project
spreadsheet = project.spreadsheet_import
#options = project.spreadsheet_options
layout = self.layout
row = layout.row(align=True)
row.label(text='Source')
#row.alignment='RIGHT'
row.operator("vse_toolbox.spreadsheet_from_clipboard", text='Clipboard', icon='PASTEDOWN')
row.operator("vse_toolbox.spreadsheet_from_file", text='File', icon='FILE').filepath = self.filepath
col = layout.column(align=False)
row = col.row(align=True, heading='Custom Asset Name')
row.use_property_split = True
row.use_property_decorate = False
row.prop(spreadsheet, 'use_custom_name', text='')
sub = row.row(align=False)
sub.enabled = spreadsheet.use_custom_name
sub.prop(spreadsheet, 'custom_name', text='')
sub.label(icon='BLANK1')
row = layout.row()
row.template_list("VSETB_UL_spreadsheet_import", "spreadsheet_import", spreadsheet, "cells", spreadsheet, "cell_index", rows=8)
col_tool = row.column(align=True)
#bpy.types.VSETB_PT_presets.draw_panel_header(col_tool)
#col_tool.operator('wm.call_menu', icon="PRESET").name = 'VSETB_MT_spreadsheet_presets'
#col_tool.operator('vse_toolbox.load_spreadsheet_preset', icon='PRESET', text="")
op = col_tool.operator('wm.call_panel', icon="PRESET", emboss=False, text='')
op.name = 'VSETB_PT_presets'
op.keep_open = False
#col_tool.separator()
#col_tool.operator('vse_toolbox.spreadsheet_move', icon='TRIA_UP', text="").direction = 'UP'
#col_tool.operator('vse_toolbox.spreadsheet_move', icon='TRIA_DOWN', text="").direction = 'DOWN'
col = layout.column()
col.use_property_split = True
col.use_property_decorate = False
#col.separator()
row = col.row(align=False)
col.prop(spreadsheet, "separator", expand=True, text='Separator')
if self.filepath.endswith('.csv'):
col.prop(spreadsheet, "delimiter", expand=True, text='Delimiter')
col.prop(spreadsheet, 'import_casting', text='Import Casting')
col.prop(spreadsheet, 'import_custom_data', text='Import Custom Data')
col.prop(spreadsheet, 'update_edit', text='Update Edit')
col.separator()
def execute(self, context):
scn = context.scene
settings = get_scene_settings()
project = settings.active_project
spreadsheet = project.spreadsheet_import
sequencer = scn.sequence_editor.sequences
assets_missing = set()
# Import Edit
nb_frames_cell = next((c for c in spreadsheet.cells if c.import_name=='Nb Frames'), None)
channel = get_channel_index('Shots')
header = SPREADSHEET[0]
#print(SPREADSHEET[:2])
cell_types = project.get_cell_types()
cell_names = {k: spreadsheet.cells[k].import_name for k in header if k}
#separator = spreadsheet.separator.replace('\\n', '\n').replace('\\t', '\t').replace('\\r', '\r')
shot_strips = get_strips('Shots')
frame_start = scn.frame_start
for row in SPREADSHEET[1:]:
#print(row)
cell_data = {cell_names[k]: v for k, v in zip(header, row) if k}
#cell_data = {k: v for k, v in zip(header, row)}
shot_name = cell_data['Shot']
#strip = next((s for s in sequencer if s.vsetb_strip_settings.source_name == shot_name), None)
strip = next((s for s in shot_strips if s.name == shot_name), None)
#print(cell_data)
#print("shot_name", shot_name, strip)
if spreadsheet.update_edit and (nb_frames_cell and nb_frames_cell.enabled):
nb_frames = int(cell_data['Nb Frames'])
#print('Import Edit')
#print(frame_start, nb_frames, type(nb_frames))
frame_end = frame_start + nb_frames
if strip:
if frame_start != strip.frame_final_start or frame_end !=strip.frame_final_end:
print(f'The strip {strip.name} is updated with new range')
strip.frame_final_start = frame_start
strip.frame_final_end = frame_end
else:
strip = sequencer.new_effect(
name=shot_name,
type='COLOR',
channel=channel,
frame_start=frame_start,
frame_end=frame_end,
)
strip.blend_alpha = 0.0
strip.select = False
strip.vsetb_strip_settings.source_name = shot_name
frame_start += nb_frames
if not strip:
self.report({"WARNING"}, f'No strip found with source name {shot_name}, update the edit')
continue
strip_settings = strip.vsetb_strip_settings
if spreadsheet.import_casting:
strip_settings.casting.clear()
strip_settings.casting.update()
#print('Clear Casting of strip', strip)
for cell_name, cell_value in zip(header, row):
if not cell_name:
continue
cell = spreadsheet.cells[cell_name]
if not cell.enabled:
continue
cell_type = cell_types[cell.import_name]
if cell.import_name == 'Description' and spreadsheet.import_custom_data:
strip_settings.description = cell_value
elif cell_type == 'METADATA' and spreadsheet.import_custom_data:
metadata = project.metadata_types[cell.import_name]
setattr(strip_settings.metadata, metadata.field_name, cell_value)
elif cell_type == 'TASK_TYPE':
task_type = project.task_types[cell.import_name]
task = getattr(strip_settings.tasks, norm_str(task_type.name))
task.comment = cell_value
if cell_type == 'ASSET_TYPE' and spreadsheet.import_casting:
#print(cell_value)
asset_names = cell_value.split('\n')
# Clear the list of assets
#for asset_casting in list(strip_settings.casting):
# if asset_casting.asset.asset_type == cell_name:
# strip_settings.casting.remove(asset_casting)
for asset_name in asset_names:
if not asset_name:
continue
#print(norm_str(asset_name), norm_str(project.assets[0].tracker_name))
norm_asset_name = norm_str(asset_name)
if spreadsheet.use_custom_name:
asset = next((a for a in project.assets if norm_str(a.get('metadata', {}).get(spreadsheet.custom_name)) == norm_asset_name), None)
else:
asset = next((a for a in project.assets if norm_str(a.tracker_name) == norm_asset_name), None)
if asset:
item = next((item for item in strip_settings.casting if item.asset == asset), None)
if item:
item.instance += 1
else:
item = strip_settings.casting.add()
item.name = asset.name
item.id = asset.id
item['_name'] = asset.label
strip_settings.casting.update()
else:
assets_missing.add(f'{cell_name} / {asset_name}')
#print(f'Asset {asset_name} not found in Project for strip {strip.name}')
#self.report({'WARNING'}, f'Asset {asset_name} not found in Project for strip {strip.name}')
if assets_missing:
print('Some assets were missing')
for asset in sorted(assets_missing):
print(asset)
self.report({'WARNING'}, f'Some assets were missing {list(assets_missing)[:5]}...')
return {"FINISHED"}
class VSETB_OT_export_spreadsheet(Operator):
bl_idname = "vse_toolbox.export_spreadsheet"
bl_label = "Export Spreadsheet"
bl_description = "Export Shot data in a table as a csv or an xlsl"
bl_options = {"REGISTER", "UNDO"}
@classmethod
def poll(cls, context):
settings = get_scene_settings()
return settings.active_project
def invoke(self, context, event):
settings = get_scene_settings()
project = settings.active_project
return context.window_manager.invoke_props_dialog(self, width=375)
def draw(self, context):
scn = context.scene
settings = get_scene_settings()
project = settings.active_project
spreadsheet = project.spreadsheet_export
#options = project.spreadsheet_options
layout = self.layout
row = layout.row(align=True)
row.prop(spreadsheet, 'use_custom_cells', text='Custom Cells')
col = layout.column(align=False)
col.enabled = spreadsheet.use_custom_cells
row = col.row(align=True, heading='Custom Asset Name')
row.use_property_split = True
row.use_property_decorate = False
row.prop(spreadsheet, 'use_custom_name', text='')
sub = row.row(align=False)
sub.enabled = spreadsheet.use_custom_name
sub.prop(spreadsheet, 'custom_name', text='')
sub.label(icon='BLANK1')
row = col.row()
row.template_list("VSETB_UL_spreadsheet_export", "spreadsheet_export", spreadsheet, "cells", spreadsheet, "cell_index", rows=8)
col_tool = row.column(align=True)
#bpy.types.VSETB_PT_presets.draw_panel_header(col_tool)
#col_tool.operator('wm.call_menu', icon="PRESET").name = 'VSETB_MT_spreadsheet_presets'
#col_tool.operator('vse_toolbox.load_spreadsheet_preset', icon='PRESET', text="")
op = col_tool.operator('wm.call_panel', icon="PRESET", emboss=False, text='')
op.name = 'VSETB_PT_presets'
op.keep_open = False
col_tool.separator()
col_tool.operator('vse_toolbox.spreadsheet_move', icon='TRIA_UP', text="").direction = 'UP'
col_tool.operator('vse_toolbox.spreadsheet_move', icon='TRIA_DOWN', text="").direction = 'DOWN'
col = layout.column()
col.use_property_split = True
col.use_property_decorate = False
col.separator()
row = col.row(align=False)
row.prop(spreadsheet, "format", expand=True, text='Format')
row.prop(spreadsheet, 'show_settings', text='', icon='PREFERENCES')
if spreadsheet.show_settings:
col.prop(spreadsheet, "separator", expand=True, text='Separator')
if spreadsheet.format == 'csv':
col.prop(spreadsheet, "delimiter", expand=True, text='Delimiter')
if spreadsheet.format != 'Clipboard':
col.separator()
col.prop(spreadsheet, 'open_folder', text='Open Folder')
col.prop(spreadsheet, 'export_path', text='Export Path')
def execute(self, context):
#self.report({'ERROR'}, f'Export not implemented yet.')
prefs = get_addon_prefs()
settings = get_scene_settings()
project = settings.active_project
spreadsheet = project.spreadsheet_export
episode = settings.active_episode
rows = []
# Header
if spreadsheet.use_custom_cells:
cells = [cell for cell in spreadsheet.cells if cell.enabled]
rows.append([cell.export_name for cell in cells])
else:
cells = spreadsheet.cells
rows.append([cell.name for cell in cells])
print(rows)
#raise Exception('')
separator = spreadsheet.separator.replace('\\n', '\n').replace('\\t', '\t').replace('\\r', '\r')
delimiter = spreadsheet.delimiter.replace('\\n', '\n').replace('\\t', '\t').replace('\\r', '\r')
for strip in get_strips('Shots'):
row = []
for cell in cells:
#print(cell.field_name)
if cell.type == "METADATA":
row += [getattr(strip.vsetb_strip_settings.metadata, cell.field_name)]
elif cell.type == "ASSET_TYPE":
asset_castings = []
for asset_casting in strip.vsetb_strip_settings.casting:
asset = asset_casting.asset
if not asset:
self.report({"WARNING"}, f"The asset {asset_casting['_name']} is missing on strip {strip.name}")
continue
if not asset.asset_type == cell.name:
continue
if spreadsheet.use_custom_name and spreadsheet.use_custom_cells:
if asset.get('metadata', {}).get(spreadsheet.custom_name):
asset_castings += [asset['metadata'][spreadsheet.custom_name]]*asset_casting.instance
else:
self.report({'ERROR'}, f'The asset {asset.tracker_name} has no data {spreadsheet.custom_name}')
else:
asset_castings += [asset.tracker_name]*asset_casting.instance
row += [separator.join(asset_castings)]
elif cell.type == 'TASK_TYPE':
row += [getattr(strip.vsetb_strip_settings.tasks, norm_str(cell.name)).comment]
elif cell.field_name == 'EPISODE':
row += [settings.active_episode.name]
elif cell.field_name == 'SEQUENCE':
row += [get_strip_sequence_name(strip)]
elif cell.field_name == 'SHOT':
row += [strip.name]
elif cell.field_name == 'DESCRIPTION':
row += [strip.vsetb_strip_settings.description]
elif cell.field_name == 'NB_FRAMES':
row += [strip.frame_final_duration]
rows.append(row)
#print(rows)
export_path = Path(os.path.abspath(bpy.path.abspath(spreadsheet.export_path)))
export_name = export_path.name
if export_path.suffix or export_name.endswith('{ext}'):
export_path = export_path.parent
else: # It's a directory
if project.type == 'TVSHOW':
export_name = '{date}_{project}_{episode}_{tracker}_shots.{ext}'
else:
export_name = '{date}_{project}_{tracker}_shots.{ext}'
date = datetime.now().strftime('%Y_%m_%d')
project_name = project.name.replace(' ', '_').lower()
episode_name = episode.name.replace(' ', '_').lower() if episode else 'episode'
ext = spreadsheet.format.lower()
export_name = export_name.format(date=date, project=project_name,
episode=episode_name, tracker=settings.tracker_name.lower(), ext=ext)
export_path = export_path / export_name
#2023_04_11_kitsu_boris_ep01_shots
export_path.parent.mkdir(parents=True, exist_ok=True)
if export_path.exists():
try:
print('Removing File', export_path)
export_path.unlink()
except Exception:
self.report({'ERROR'}, f'You need to close the file {export_path}')
if spreadsheet.format == 'csv':
print('Writing .csv file to', export_path)
with open(str(export_path), 'w', newline='\n', encoding='utf-8') as f:
writer = csv.writer(f, delimiter=spreadsheet.delimiter)
for row in rows:
#print(row)
writer.writerow(row)
elif spreadsheet.format == 'xlsx':
try:
import openpyxl
except ModuleNotFoundError:
self.report({'INFO'}, 'Installing openpyxl...')
openpyxl = install_module('openpyxl')
from openpyxl import Workbook
workbook = Workbook()
worksheet = workbook.active
workbook.active.title = 'BKL'
for row in rows:
worksheet.append(row)
for col in worksheet.columns:
letter = col[0].column_letter
worksheet.column_dimensions[letter].auto_size = True
# Save the file
workbook.save(str(export_path))
elif spreadsheet.format == 'xls':
try:
import xlwt
except ModuleNotFoundError:
self.report({'INFO'}, 'Installing xlwt...')
xlwt = install_module('xlwt')
workbook = xlwt.Workbook()
worksheet = workbook.add_sheet('BKL')
for row_index, row in enumerate(rows):
for col_index, cell_value in enumerate(row):
worksheet.write(row_index, col_index, cell_value)
workbook.save(str(export_path))
elif spreadsheet.format == 'Clipboard':
csv_buffer = StringIO()
# Write CSV data to the StringIO buffer
csv_writer = csv.writer(csv_buffer, delimiter='\t')
csv_writer.writerows(rows)
context.window_manager.clipboard = csv_buffer.getvalue()
if spreadsheet.open_folder:
open_file(export_path, select=True)
return {"FINISHED"}
classes = (
VSETB_OT_add_export_spreadsheet_preset,
VSETB_MT_export_spreadsheet_presets,
VSETB_OT_spreadsheet_from_file,
VSETB_OT_spreadsheet_from_clipboard,
VSETB_OT_spreadsheet_to_clipboard,
VSETB_OT_spreadsheet_cell_move,
VSETB_OT_export_spreadsheet,
VSETB_OT_import_spreadsheet
)
def register():
for cls in classes:
bpy.utils.register_class(cls)
def unregister():
for cls in reversed(classes):
bpy.utils.unregister_class(cls)