Add optional close tab warning / Fix base64 encoding for python2

Author: luckylyk

dwpicker/main.py

@@ -24,8 +24,8 @@
     AUTO_FOCUS_BEHAVIOR, AUTO_SWITCH_TAB, CHECK_IMAGES_PATHS,
     DISABLE_IMPORT_CALLBACKS, DISPLAY_QUICK_OPTIONS, INSERT_TAB_AFTER_CURRENT,
     LAST_OPEN_DIRECTORY, LAST_IMPORT_DIRECTORY, LAST_SAVE_DIRECTORY,
-    NAMESPACE_TOOLBAR, USE_ICON_FOR_UNSAVED_TAB, save_optionvar,
-    append_recent_filename, save_opened_filenames)
+    NAMESPACE_TOOLBAR, USE_ICON_FOR_UNSAVED_TAB, WARN_ON_TAB_CLOSED,
+    save_optionvar, append_recent_filename, save_opened_filenames)
 from dwpicker.picker import PickerView, list_targets
 from dwpicker.preference import PreferencesWindow
 from dwpicker.qtutils import set_shortcut, icon, maya_main_window, DockableBase
@@ -62,6 +62,9 @@
 WINDOW_TITLE = "DreamWall - Picker"
 WINDOW_CONTROL_NAME = "dwPickerWindow"
 CLOSE_CALLBACK_COMMAND = "import dwpicker;dwpicker._dwpicker.close_event()"
+CLOSE_TAB_WARNING = """\
+Close the tab will remove completely the picker data from the scene.
+Are you sure to continue ?"""
 
 
 def build_multiple_shapes(targets, override):
@@ -368,11 +371,11 @@ def clear(self):
             self.close_tab(i, force=True)
 
     def close_tab(self, index, force=False, store=False):
-        conditions = (
-            self.modified_states[index]
-            and force is False
-            and not self.save_tab(index))
-        if conditions:
+        if self.modified_states[index] and force is False:
+            if not self.save_tab(index):
+                return
+        elif (cmds.optionVar(query=WARN_ON_TAB_CLOSED) and
+              not question('Warning', CLOSE_TAB_WARNING)):
             return
 
         editor = self.editors.pop(index)
@@ -469,7 +472,7 @@ def call_save(self, index=None):
         filename = self.filenames[index]
         if not filename:
             return self.call_save_as(index=index)
-        self.save_picker(index, filename)
+        return self.save_picker(index, filename)
 
     def call_save_as(self, index=None):
         index = self.tab.currentIndex() if type(index) is not int else index

dwpicker/optionvar.py

@@ -37,6 +37,7 @@
 USE_BASE64_DATA_ENCODING = 'dwpicker_use_base64_data_encoding'
 USE_ICON_FOR_UNSAVED_TAB = 'dwpicker_use_icon_for_unsaved_tab'
 ZOOM_BUTTON = 'dwpicker_picker_zoom_mouse_button'
+WARN_ON_TAB_CLOSED = 'dwpicker_warn_on_tab_closed'
 ZOOM_SENSITIVITY = 'dwpicker_zoom_sensitivity'
 
 
@@ -70,6 +71,7 @@
     TRIGGER_REPLACE_ON_MIRROR: 0,
     USE_BASE64_DATA_ENCODING: 0,
     USE_ICON_FOR_UNSAVED_TAB: 1,
+    WARN_ON_TAB_CLOSED: 0,
     ZOOM_BUTTON: ZOOM_BUTTONS[2],
     ZOOM_SENSITIVITY: 50
 }

dwpicker/preference.py

@@ -6,7 +6,8 @@
     CHECK_IMAGES_PATHS, DISPLAY_QUICK_OPTIONS, DISABLE_IMPORT_CALLBACKS,
     INSERT_TAB_AFTER_CURRENT, NAMESPACE_TOOLBAR, SYNCHRONYZE_SELECTION,
     TRIGGER_REPLACE_ON_MIRROR, USE_BASE64_DATA_ENCODING,
-    USE_ICON_FOR_UNSAVED_TAB, ZOOM_SENSITIVITY, ZOOM_BUTTON, ZOOM_BUTTONS)
+    USE_ICON_FOR_UNSAVED_TAB, WARN_ON_TAB_CLOSED, ZOOM_SENSITIVITY,
+    ZOOM_BUTTON, ZOOM_BUTTONS)
 
 
 MAX_SENSITIVITY = 500
@@ -37,6 +38,8 @@ def __init__(self, callback=None, parent=None):
         self.unsaved_tab_icon = QtWidgets.QCheckBox(text)
         text = "Insert new tab after current tab."
         self.insert_after_current = QtWidgets.QCheckBox(text)
+        text = "Warning before closing a tab."
+        self.warn_on_tab_close = QtWidgets.QCheckBox(text)
         self.ui_group = QtWidgets.QGroupBox("Ui")
         self.ui_layout = QtWidgets.QVBoxLayout(self.ui_group)
         self.ui_layout.addWidget(self.namespace_toolbar)
@@ -47,6 +50,7 @@ def __init__(self, callback=None, parent=None):
         self.ui_layout.addWidget(self.check_images_paths)
         self.ui_layout.addWidget(self.unsaved_tab_icon)
         self.ui_layout.addWidget(self.insert_after_current)
+        self.ui_layout.addWidget(self.warn_on_tab_close)
 
         text = "Encode in-scene data as base64."
         self.use_base64_encoding = QtWidgets.QCheckBox(text)
@@ -101,6 +105,7 @@ def __init__(self, callback=None, parent=None):
         self.unsaved_tab_icon.released.connect(self.save_ui_states)
         self.sychronize.released.connect(self.save_ui_states)
         self.search_on_mirror.released.connect(self.save_ui_states)
+        self.warn_on_tab_close.released.connect(self.save_ui_states)
         self.zoom_sensitivity.valueChanged.connect(self.save_ui_states)
         self.zoom_button.currentIndexChanged.connect(self.save_ui_states)
 
@@ -121,6 +126,8 @@ def load_ui_states(self):
         self.use_base64_encoding.setChecked(state)
         state = bool(cmds.optionVar(query=USE_ICON_FOR_UNSAVED_TAB))
         self.unsaved_tab_icon.setChecked(state)
+        state = bool(cmds.optionVar(query=WARN_ON_TAB_CLOSED))
+        self.warn_on_tab_close.setChecked(state)
         state = bool(cmds.optionVar(query=INSERT_TAB_AFTER_CURRENT))
         self.insert_after_current.setChecked(state)
         value = cmds.optionVar(query=AUTO_FOCUS_BEHAVIOR)
@@ -155,6 +162,8 @@ def save_ui_states(self, *_):
         save_optionvar(AUTO_FOCUS_BEHAVIOR, value)
         value = int(self.search_on_mirror.isChecked())
         save_optionvar(TRIGGER_REPLACE_ON_MIRROR, value)
+        value = int(self.warn_on_tab_close.isChecked())
+        save_optionvar(WARN_ON_TAB_CLOSED, value)
         save_optionvar(ZOOM_BUTTON, self.zoom_button.currentText())
         value = MAX_SENSITIVITY - int(self.zoom_sensitivity.value()) + 1
         save_optionvar(ZOOM_SENSITIVITY, value)

dwpicker/scenedata.py

@@ -1,3 +1,4 @@
+import sys
 import json
 import base64
 
@@ -28,9 +29,7 @@ def create_picker_holder_node():
 
 
 def store_local_picker_data(pickers):
-    data = json.dumps(pickers)
-    if cmds.optionVar(query=USE_BASE64_DATA_ENCODING):
-        data = base64.b64encode(bytes(data, "utf-8"))
+    data = encode_data(pickers)
     node = get_picker_holder_node()
     cmds.setAttr(node + '.' + PICKER_HOLDER_ATTRIBUTE, data, type='string')
     clean_stray_picker_holder_nodes()
@@ -46,6 +45,16 @@ def load_local_picker_data():
     return pickers
 
 
+def encode_data(pickers):
+    data = json.dumps(pickers)
+    if not cmds.optionVar(query=USE_BASE64_DATA_ENCODING):
+        return data
+    # Ensure backward compatibility.
+    if sys.version_info[0] == 2:
+        return base64.b64encode(bytes(data))
+    return base64.b64encode(bytes(data, "utf-8"))
+
+
 def decode_data(data):
     try:
         return json.loads(data)