Add storage implementation for referencing (wip)

Author: kalemas

dwpicker/scenedata_dag.py

@@ -6,6 +6,7 @@
 imp.reload(dwpicker)
 import dwpicker.scenedata
 from maya import cmds
+from dwpicker.namespace import maya_namespace
 
 
 def shortname(fullname):
@@ -121,10 +122,84 @@ def store(self, pickers):
     def cleanup(self):
         pass
 
-#data = timeit(dwpicker.scenedata.DefaultSceneStorage().load)()
-#data = timeit(SceneDagStorage().load)()
-#SceneDagStorage().store(data)
-#dwpicker.scenedata.DefaultSceneStorage().store(data)
-#pprint.pprint(data)
-#dwpicker.show(storage_class=SceneDagStorage)
-#dwpicker.show(storage_class=dwpicker.scenedata.DefaultSceneStorage)
+
+def node_namespace(name):
+    # TODO dwpicker.namespace.node_namespace do not handle
+    # nested namespaces
+    return (name.rsplit(':', 1)[:-1] or [None])[-1]
+
+
+class SceneDictDiffStorage:
+    PICKER_HOLDER_NODE = 'DwPicker'
+    PICKER_HOLDER_ATTRIBUTE = '_dwpicker_data'
+    # LS_EXP = ["*." + PICKER_HOLDER_ATTRIBUTE, "*:*." + PICKER_HOLDER_ATTRIBUTE]
+    LS_EXP = [PICKER_HOLDER_NODE, "*:" + PICKER_HOLDER_NODE]
+
+    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('dagContainer', name=self.PICKER_HOLDER_NODE, skipSelect=True)
+        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 = []
+        added_pickers = set()
+        # start from root holders
+        for node in sorted(nodes, key=lambda x: len(x)):
+            for picker_node in cmds.listRelatives(node, children=True) or []:
+                if picker_node in added_pickers:
+                    continue
+                namespace = node_namespace(picker_node)
+                data = cmds.getAttr(picker_node + '.' + self.PICKER_HOLDER_ATTRIBUTE)
+                data = dwpicker.scenedata.decode_data(data)
+                data = dwpicker.scenedata.ensure_retro_compatibility(data)
+                if namespace:
+                    for s in data['shapes']:
+                        s['action.targets'] = ['{}:{}'.format(namespace, n)
+                                               for n in s['action.targets']]
+                data['source_node'] = picker_node
+                pickers.append(data)
+                added_pickers.add(data['source_node'])
+        return pickers
+
+    def store(self, pickers):
+        node = self._get_picker_holder_node()
+        for picker_node in cmds.listRelatives(node, children=True) or []:
+            cmds.delete(picker_node)
+        for p in pickers:
+            picker_node = p.get('source_node') or p.get('general', {}).get('name') or 'Picker'
+            picker_node = cmds.createNode('dagContainer', name=picker_node, parent=node, skipSelect=True)
+            if p.get('source_node') == picker_node:
+                del p['source_node']
+            cmds.addAttr(
+                picker_node, longName=self.PICKER_HOLDER_ATTRIBUTE, dataType='string')
+            data = dwpicker.scenedata.encode_data(p)
+            cmds.setAttr(picker_node + '.' + self.PICKER_HOLDER_ATTRIBUTE, data, type='string')
+
+    def cleanup(self):
+        pass
+
+
+# data = timeit(dwpicker.scenedata.DefaultSceneStorage().load)()
+# data = timeit(SceneDagStorage().load)()
+# SceneDagStorage().store(data)
+# dwpicker.scenedata.DefaultSceneStorage().store(data)
+# pprint.pprint(data)
+# dwpicker.show(storage_class=SceneDagStorage)
+# dwpicker.show(storage_class=dwpicker.scenedata.DefaultSceneStorage)
+# import dwpicker.scenedata_dag
+# dwpicker.show(storage_class=dwpicker.scenedata_dag.SceneDictDiffStorage)