Refactor storage

dwpicker/__init__.py

@@ -8,12 +8,14 @@
 _dwpicker = None
 
 
-def show(editable=True, pickers=None, ignore_scene_pickers=False):
+def show(
+        editable=True, pickers=None, ignore_scene_pickers=False, 
+        storage_class=None):
     ensure_optionvars_exists()
     global _dwpicker
     if not _dwpicker:
         warn_if_update_available()
-        _dwpicker = DwPicker()
+        _dwpicker = DwPicker(storage_class=storage_class)
 
     try:
         _dwpicker.show(dockable=True)

dwpicker/main.py

@@ -31,9 +31,7 @@
 from dwpicker.qtutils import set_shortcut, icon, maya_main_window, DockableBase
 from dwpicker.quick import QuickOptions
 from dwpicker.references import ensure_images_path_exists
-from dwpicker.scenedata import (
-    load_local_picker_data, store_local_picker_data,
-    clean_stray_picker_holder_nodes)
+from dwpicker.scenedata import DefaultSceneStorage
 from dwpicker.templates import BUTTON, PICKER, BACKGROUND
 from dwpicker.undo import UndoManager
 
@@ -77,7 +75,7 @@ def build_multiple_shapes(targets, override):
 
 
 class DwPicker(DockableBase, QtWidgets.QWidget):
-    def __init__(self):
+    def __init__(self, storage_class=None):
         super(DwPicker, self).__init__(control_name=WINDOW_CONTROL_NAME)
         self.setWindowTitle(WINDOW_TITLE)
         shortcut_context = QtCore.Qt.WidgetWithChildrenShortcut
@@ -99,6 +97,8 @@ def __init__(self):
         self.pickers = []
         self.filenames = []
         self.modified_states = []
+        self.storage = (storage_class or DefaultSceneStorage)()
+
         self.preferences_window = PreferencesWindow(
             callback=self.load_ui_states, parent=maya_main_window())
         self.preferences_window.disable_import_callbacks.released.connect(
@@ -324,23 +324,23 @@ def auto_switch_tab(self, *_, **__):
 
     def load_saved_pickers(self, *_, **__):
         self.clear()
-        pickers = load_local_picker_data()
+        pickers = self.storage.load()
         if cmds.optionVar(query=CHECK_IMAGES_PATHS):
             picker = ensure_images_path_exists(pickers)
         for picker in pickers:
             self.add_picker(picker)
-        clean_stray_picker_holder_nodes()
+        self.storage.cleanup()
 
     def store_local_pickers_data(self):
         if not self.editable:
             return
 
         if not self.tab.count():
-            store_local_picker_data([])
+            self.storage.store([])
             return
 
         pickers = [self.picker_data(i) for i in range(self.tab.count())]
-        store_local_picker_data(pickers)
+        self.storage.store(pickers)
 
     def save_tab(self, index):
         msg = (

dwpicker/scenedata.py

@@ -9,42 +9,6 @@
 from dwpicker.namespace import maya_namespace
 
 
-PICKER_HOLDER_NODE = '_dwpicker_data'
-PICKER_HOLDER_ATTRIBUTE = '_dwpicker_data'
-LS_EXP = ["*." + PICKER_HOLDER_ATTRIBUTE, "*:*." + PICKER_HOLDER_ATTRIBUTE]
-
-
-def get_picker_holder_node():
-    if cmds.objExists(PICKER_HOLDER_NODE):
-        return PICKER_HOLDER_NODE
-    return create_picker_holder_node()
-
-
-def create_picker_holder_node():
-    with maya_namespace(":"):
-        node = cmds.createNode('script', name=PICKER_HOLDER_NODE)
-    cmds.setAttr(node + '.nodeState', 1)
-    cmds.addAttr(node, longName=PICKER_HOLDER_ATTRIBUTE, dataType='string')
-    return node
-
-
-def store_local_picker_data(pickers):
-    data = encode_data(pickers)
-    node = get_picker_holder_node()
-    cmds.setAttr(node + '.' + PICKER_HOLDER_ATTRIBUTE, data, type='string')
-    clean_stray_picker_holder_nodes()
-
-
-def load_local_picker_data():
-    nodes = list_picker_holder_nodes()
-    pickers = []
-    for node in nodes:
-        data = cmds.getAttr(node + '.' + PICKER_HOLDER_ATTRIBUTE)
-        data = decode_data(data)
-        pickers.extend(ensure_retro_compatibility(p) for p in data)
-    return pickers
-
-
 def encode_data(pickers):
     data = json.dumps(pickers)
     if not cmds.optionVar(query=USE_BASE64_DATA_ENCODING):
@@ -62,28 +26,63 @@ def decode_data(data):
         return json.loads(base64.b64decode(data))
 
 
-def list_picker_holder_nodes():
-    """
-    Look up in the scene all the nodes holding an attribute named
-    "_dwpicker_holder" which are not set on the "_dwpicker_holder" node.
-    This mignt happed if a node node is imported (creating a namespace or a
-    incrementation).
-    """
-    return [node.split(".")[0] for node in cmds.ls(LS_EXP)]
-
-
-def clean_stray_picker_holder_nodes():
-    """
-    If the scene contains multiple picker holder nodes, we remove them
-    automatically to avoid repeated pickers.
-    """
-    for node in list_picker_holder_nodes():
-        if node == PICKER_HOLDER_NODE:
-            continue
-        try:
-            cmds.delete(node)
-        except:
-            # Node is locked or in reference and cannot be removed.
-            # As we cant remove it, we reset his data to avoid double pickers.
-            cmds.setAttr(
-                node + "." + PICKER_HOLDER_ATTRIBUTE, "", dataType="string")
+class DefaultSceneStorage:
+    PICKER_HOLDER_NODE = '_dwpicker_data'
+    PICKER_HOLDER_ATTRIBUTE = '_dwpicker_data'
+    LS_EXP = ["*." + PICKER_HOLDER_ATTRIBUTE, "*:*." + PICKER_HOLDER_ATTRIBUTE]
+
+    def _get_picker_holder_node(self):
+        if cmds.objExists(self.PICKER_HOLDER_NODE):
+            return self.PICKER_HOLDER_NODE
+        return self._create_picker_holder_node()
+
+    def _create_picker_holder_node(self):
+        with maya_namespace(":"):
+            node = cmds.createNode('script', name=self.PICKER_HOLDER_NODE)
+        cmds.setAttr(node + '.nodeState', 1)
+        cmds.addAttr(
+            node, longName=self.PICKER_HOLDER_ATTRIBUTE, dataType='string')
+        return node
+
+    def _list_picker_holder_nodes(self):
+        """
+        Look up in the scene all the nodes holding an attribute named
+        "_dwpicker_holder" which are not set on the "_dwpicker_holder" node.
+        This mignt happed if a node node is imported (creating a namespace or a
+        incrementation).
+        """
+        return [node.split(".")[0] for node in cmds.ls(self.LS_EXP)]
+
+    def load(self):
+        nodes = self._list_picker_holder_nodes()
+        pickers = []
+        for node in nodes:
+            data = cmds.getAttr(node + '.' + self.PICKER_HOLDER_ATTRIBUTE)
+            data = decode_data(data)
+            pickers.extend(ensure_retro_compatibility(p) for p in data)
+        return pickers
+
+    def store(self, pickers):
+        data = encode_data(pickers)
+        node = self._get_picker_holder_node()
+        cmds.setAttr(
+            node + '.' + self.PICKER_HOLDER_ATTRIBUTE, data, type='string')
+        self.cleanup()
+
+    def cleanup(self):
+        """
+        If the scene contains multiple picker holder nodes, we remove them
+        automatically to avoid repeated pickers.
+        """
+        for node in self._list_picker_holder_nodes():
+            if node == self.PICKER_HOLDER_NODE:
+                continue
+            try:
+                cmds.delete(node)
+            except:
+                # Node is locked or in reference and cannot be removed.
+                # As we cant remove it, we reset his data to avoid double 
+                # pickers.
+                cmds.setAttr(
+                    node + "." + self.PICKER_HOLDER_ATTRIBUTE, "", 
+                    dataType="string")