For #10692: Option to use add/remove/set modes when updating multi-entity fields (#111)

victoriagrey · victoriagrey · commit c89da308edff · 2016-04-25T12:56:06.000-07:00
This update will allow for updating multi-entity fields without having to pull down the field's entire dataset first.

It uses the new parameter multi_entity_update_modes on the update method to specify one of 'add', 'remove', or 'set'.
diff --git a/shotgun_api3/shotgun.py b/shotgun_api3/shotgun.py
@@ -843,7 +843,7 @@ def create(self, entity_type, data, return_fields=None):
 
         return result
 
-    def update(self, entity_type, entity_id, data):
+    def update(self, entity_type, entity_id, data, multi_entity_update_modes=None):
         """Updates the specified entity with the supplied data.
 
         :param entity_type: Required, entity type (string) to update.
@@ -852,6 +852,16 @@ def update(self, entity_type, entity_id, data):
 
         :param data: Required, dict fields to update on the entity.
 
+        :param multi_entity_update_modes: Optional, dict of what update mode to
+        use when updating a multi-entity link field.  The keys in the dict are
+        the fields to set the mode for and the values from the dict are one
+        of "set", "add", or "remove". The default behavior if mode is not
+        specified for a field is 'set'. For example, on the 'Sequence' entity,
+        to append to the 'shots' field and remove from the 'assets' field, you
+        would specify:
+
+            multi_entity_update_modes={"shots":"add", "assets":"remove"}
+
         :returns: dict of the fields updated, with the entity_type and
         id added.
         """
@@ -871,7 +881,10 @@ def update(self, entity_type, entity_id, data):
             params = {
                 "type" : entity_type,
                 "id" : entity_id,
-                "fields" : self._dict_to_list(data)
+                "fields" : self._dict_to_list(
+                    data,
+                    extra_data=self._dict_to_extra_data(
+                        multi_entity_update_modes, "multi_entity_update_mode"))
             }
             record = self._call_rpc("update", params)
             result = self._parse_records(record)[0]
@@ -941,7 +954,7 @@ def batch(self, requests):
         :param requests: A list of dict's of the form which have a
             request_type key and also specifies:
             - create: entity_type, data dict of fields to set
-            - update: entity_type, entity_id, data dict of fields to set
+            - update: entity_type, entity_id, data dict of fields to set, optionally multi_entity_update_modes
             - delete: entity_type and entity_id
 
         :returns: A list of values for each operation, create and update
@@ -982,7 +995,12 @@ def _required_keys(message, required_keys, data):
                                ['entity_id', 'data'],
                                req)
                 request_params['id'] = req['entity_id']
-                request_params['fields'] = self._dict_to_list(req["data"])
+                request_params['fields'] = self._dict_to_list(req["data"],
+                    extra_data=self._dict_to_extra_data(
+                        req.get("multi_entity_update_modes"),
+                        "multi_entity_update_mode"))
+                if "multi_entity_update_mode" in req:
+                    request_params['multi_entity_update_mode'] = req["multi_entity_update_mode"]
             elif req["request_type"] == "delete":
                 _required_keys("Batched delete request", ['entity_id'], req)
                 request_params['id'] = req['entity_id']
@@ -2573,18 +2591,30 @@ def _build_thumb_url(self, entity_type, entity_id):
         # Comments in prev version said we can get this sometimes.
         raise RuntimeError("Unknown code %s %s" % (code, thumb_url))
 
-    def _dict_to_list(self, d, key_name="field_name", value_name="value"):
+    def _dict_to_list(self, d, key_name="field_name", value_name="value", extra_data=None):
         """Utility function to convert a dict into a list dicts using the
         key_name and value_name keys.
 
-        e.g. d {'foo' : 'bar'} changed to [{'field_name':'foo, 'value':'bar'}]
-        """
-
-        return [
-            {key_name : k, value_name : v }
-            for k, v in (d or {}).iteritems()
-        ]
+        e.g. d {'foo' : 'bar'} changed to [{'field_name':'foo', 'value':'bar'}]
 
+        Any dictionary passed in via extra_data will be merged into the resulting dictionary.
+        e.g. d as above and extra_data of {'foo': {'thing1': 'value1'}} changes into
+        [{'field_name': 'foo', 'value': 'bar', 'thing1': 'value1'}]
+        """
+        ret = []
+        for k, v in (d or {}).iteritems():
+            d = {key_name: k, value_name: v}
+            d.update((extra_data or {}).get(k, {}))
+            ret.append(d)
+        return ret
+
+    def _dict_to_extra_data(self, d, key_name="value"):
+        """Utility function to convert a dict into a dict compatible with the extra_data arg
+        of _dict_to_list
+
+        e.g. d {'foo' : 'bar'} changed to {'foo': {"value": 'bar'}]
+        """
+        return dict([(k, {key_name: v}) for (k,v) in (d or {}).iteritems()])
 
 # Helpers from the previous API, left as is.
 
diff --git a/tests/base.py b/tests/base.py
@@ -26,6 +26,7 @@ def __init__(self, *args, **kws):
         self.asset          = None
         self.version        = None
         self.note           = None
+        self.playlist       = None
         self.task           = None
         self.ticket         = None
         self.human_password = None
@@ -180,9 +181,12 @@ def _setup_mock_data(self):
         self.version    = { 'id':5,
                             'code':self.config.version_code,
                             'type':'Version' }
-        self.ticket    = { 'id':6,
+        self.ticket     = { 'id':6,
                             'title':self.config.ticket_title,
                             'type':'Ticket' }
+        self.playlist   = { 'id':7,
+                            'code':self.config.playlist_code,
+                            'type':'Playlist'}
 
 class LiveTestBase(TestBase):
     '''Test base for tests relying on connection to server.'''
@@ -236,6 +240,11 @@ def _setup_db(self, config):
                 'content':'anything'}
         self.note = _find_or_create_entity(self.sg, 'Note', data, keys)
 
+        keys = ['code','project']
+        data = {'project':self.project,
+                'code':self.config.playlist_code}
+        self.playlist = _find_or_create_entity(self.sg, 'Playlist', data, keys)
+
         keys = ['code', 'entity_type']
         data = {'code': 'wrapper test step',
                 'entity_type': 'Shot'}
@@ -262,7 +271,6 @@ def _setup_db(self, config):
                 'mac_path':'nowhere',
                 'windows_path':'nowhere',
                 'linux_path':'nowhere'}
-
         self.local_storage = _find_or_create_entity(self.sg, 'LocalStorage', data, keys)
 
 
diff --git a/tests/example_config b/tests/example_config
@@ -23,4 +23,5 @@ asset_code     : Sg unittest asset
 version_code   : Sg unittest version
 shot_code      : Sg unittest shot
 task_content   : Sg unittest task
-ticket_title    : Sg unittest ticket
+ticket_title   : Sg unittest ticket
+playlist_code  : Sg unittest playlist
diff --git a/tests/test_api.py b/tests/test_api.py
@@ -639,7 +639,7 @@ class TestDataTypes(base.LiveTestBase):
     '''Test fields representing the different data types mapped on the server side.
 
      Untested data types:  password, percent, pivot_column, serializable, image, currency
-                           multi_entity, system_task_type, timecode, url, uuid, url_template
+                           system_task_type, timecode, url, uuid, url_template
     '''
     def setUp(self):
         super(TestDataTypes, self).setUp()
@@ -739,6 +739,61 @@ def test_set_list(self):
         self.assertEqual(expected, actual)
 
 
+    def test_set_multi_entity(self):
+        sg = shotgun_api3.Shotgun( self.config.server_url,
+                                   self.config.script_name,
+                                   self.config.api_key )
+        keys = ['project','user','code']
+        data = {'project':self.project,
+                'user':self.human_user,
+                'code':'Alpha'}
+        version_1 = base._find_or_create_entity(sg, 'Version', data, keys)
+        data = {'project':self.project,
+                'user':self.human_user,
+                'code':'Beta'}
+        version_2 = base._find_or_create_entity(sg, 'Version', data, keys)
+
+        entity = 'Playlist'
+        entity_id = self.playlist['id']
+        field_name = 'versions'
+
+        # Default set behaviour
+        pos_values = [[version_1, version_2]]
+        expected, actual = self.assert_set_field(entity, entity_id,
+            field_name, pos_values)
+        self.assertEqual(len(expected), len(actual))
+        self.assertEqual(
+            sorted([x['id'] for x in expected]),
+            sorted([x['id'] for x in actual])
+        )
+
+        # Multi-entity remove mode
+        pos_values = [[version_1]]
+        expected, actual = self.assert_set_field(entity, entity_id,
+            field_name, pos_values, multi_entity_update_mode='remove')
+        self.assertEqual(1, len(actual))
+        self.assertEqual(len(expected), len(actual))
+        self.assertNotEqual(expected[0]['id'],actual[0]['id'])
+        self.assertEqual(version_2['id'], actual[0]['id'])
+
+        # Multi-entity add mode
+        pos_values = [[version_1]]
+        expected, actual = self.assert_set_field(entity, entity_id,
+            field_name, pos_values, multi_entity_update_mode='add')
+        self.assertEqual(2, len(actual))
+        self.assertTrue(version_1['id'] in [x['id'] for x in actual])
+
+        # Multi-entity set mode
+        pos_values = [[version_1, version_2]]
+        expected, actual = self.assert_set_field(entity, entity_id,
+            field_name, pos_values, multi_entity_update_mode='set')
+        self.assertEqual(len(expected), len(actual))
+        self.assertEqual(
+            sorted([x['id'] for x in expected]),
+            sorted([x['id'] for x in actual])
+        )
+
+
     def test_set_number(self):
         entity = 'Shot'
         entity_id = self.shot['id']
@@ -805,13 +860,17 @@ def test_set_text_html_entity(self):
                                                  pos_values)
         self.assertEqual(expected, actual)
 
-    def assert_set_field(self, entity, entity_id, field_name, pos_values):
+    def assert_set_field(self, entity, entity_id, field_name, pos_values, multi_entity_update_mode=None):
         query_result = self.sg.find_one(entity,
                                          [['id', 'is', entity_id]],
                                          [field_name])
         initial_value = query_result[field_name]
         new_value = (initial_value == pos_values[0] and pos_values[1]) or pos_values[0]
-        self.sg.update(entity, entity_id, {field_name:new_value})
+        if multi_entity_update_mode:
+            self.sg.update(entity, entity_id, {field_name:new_value},
+                multi_entity_update_modes={field_name:multi_entity_update_mode})
+        else:
+            self.sg.update(entity, entity_id, {field_name:new_value})
         new_values = self.sg.find_one(entity,
                                      [['id', 'is', entity_id]],
                                      [field_name])
