SG-37548 Add payload optimization on update method (#363)

* Include the "in" and "not_in" for payload optimization

* Use a decorator to mock environmental variable

* Add payload optimization for update action

* Change conditional order

* Extract function to be reused. Improve testing.

* Fix typo

* Fix test

* Support multi entity, Add tests

* Apply CR feedback

* Update shotgun_api3/shotgun.py

Co-authored-by: Julien Langlois <[email protected]>

---------

Co-authored-by: Julien Langlois <[email protected]>

Author: carlos-villavicencio-adsk

shotgun_api3/shotgun.py

@@ -1132,6 +1132,28 @@ def _add_project_param(self, params, project_entity):
             params["project"] = project_entity
 
         return params
+    
+    def _translate_update_params(
+        self, entity_type, entity_id, data, multi_entity_update_modes
+    ):
+        global SHOTGUN_API_ENABLE_ENTITY_OPTIMIZATION
+
+        def optimize_field(field_dict):
+            if SHOTGUN_API_ENABLE_ENTITY_OPTIMIZATION:
+                return {k: _get_type_and_id_from_value(v) for k, v in field_dict.items()}
+            return field_dict
+
+        full_fields = self._dict_to_list(
+            data,
+            extra_data=self._dict_to_extra_data(
+                multi_entity_update_modes, "multi_entity_update_mode"
+            ),
+        )
+        return {
+            "type": entity_type,
+            "id": entity_id,
+            "fields": [optimize_field(field_dict) for field_dict in full_fields],
+        }
 
     def summarize(self,
                   entity_type,
@@ -1463,14 +1485,7 @@ def update(self, entity_type, entity_id, data, multi_entity_update_modes=None):
             upload_filmstrip_image = data.pop("filmstrip_image")
 
         if data:
-            params = {
-                "type": entity_type,
-                "id": entity_id,
-                "fields": self._dict_to_list(
-                    data,
-                    extra_data=self._dict_to_extra_data(
-                        multi_entity_update_modes, "multi_entity_update_mode"))
-            }
+            params = self._translate_update_params(entity_type, entity_id, data, multi_entity_update_modes)
             record = self._call_rpc("update", params)
             result = self._parse_records(record)[0]
         else:
@@ -4485,10 +4500,7 @@ def _translate_filters_simple(sg_filter):
         and condition["relation"] in ["is", "is_not", "in", "not_in"]
         and isinstance(values[0], dict)
     ):
-        try:
-            values = [{"type": v["type"], "id": v["id"]} for v in values]
-        except KeyError:
-            pass
+        values = [_get_type_and_id_from_value(v) for v in values]
 
     condition["values"] = values
 
@@ -4500,3 +4512,19 @@ def _version_str(version):
     Convert a tuple of int's to a '.' separated str.
     """
     return ".".join(map(str, version))
+
+
+def _get_type_and_id_from_value(value):
+    """
+    For an entity dictionary, returns a new dictionary with only the type and id keys.
+    If any of these keys are not present, the original dictionary is returned.
+    """
+    try:
+        if isinstance(value, dict):
+            return {"type": value["type"], "id": value["id"]}
+        elif isinstance(value, list):
+            return [{"type": v["type"], "id": v["id"]} for v in value]    
+    except (KeyError, TypeError):
+        LOG.debug(f"Could not optimize entity value {value}")
+
+    return value

tests/test_unit.py

@@ -290,6 +290,8 @@ def test_py_version(self, mock_sys):
 
 
 class TestFilters(unittest.TestCase):
+    maxDiff = None
+
     def test_empty(self):
         expected = {
             "logical_operator": "and",
@@ -463,6 +465,28 @@ def test_related_object_entity_optimization_is(self):
         result = api.shotgun._translate_filters(filters, "all")
         self.assertEqual(result, expected)
 
+        # Now test a non-related object. The expected result should not be optimized.
+        filters = [
+            [
+                "something",
+                "is",
+                {"foo": "foo", "bar": "bar"},
+            ],
+        ]
+        expected = {
+            "logical_operator": "and",
+            "conditions": [
+                {
+                    "path": "something",
+                    "relation": "is",
+                    "values": [{'bar': 'bar', 'foo': 'foo'}],
+                }
+            ],
+        }
+        api.Shotgun("http://server_path", "script_name", "api_key", connect=False)
+        result = api.shotgun._translate_filters(filters, "all")
+        self.assertEqual(result, expected)
+
     @mock.patch.dict(os.environ, {"SHOTGUN_API_ENABLE_ENTITY_OPTIMIZATION": "1"})
     def test_related_object_entity_optimization_in(self):
         filters = [
@@ -471,7 +495,8 @@ def test_related_object_entity_optimization_in(self):
                 "in",
                 [
                     {"foo1": "foo1", "bar1": "bar1", "id": 999, "baz1": "baz1", "type": "Anything"},
-                    {"foo2": "foo2", "bar2": "bar2", "id": 998, "baz2": "baz2", "type": "Anything"}
+                    {"foo2": "foo2", "bar2": "bar2", "id": 998, "baz2": "baz2", "type": "Anything"},
+                    {"foo3": "foo3", "bar3": "bar3"},
                 ],
             ],
         ]
@@ -489,6 +514,10 @@ def test_related_object_entity_optimization_in(self):
                         {
                             "id": 998,
                             "type": "Anything",
+                        },
+                        {
+                            "foo3": "foo3",
+                            "bar3": "bar3",
                         }
                     ],
                 }
@@ -498,6 +527,124 @@ def test_related_object_entity_optimization_in(self):
         result = api.shotgun._translate_filters(filters, "all")
         self.assertEqual(result, expected)
 
+    def test_related_object_update_entity(self):
+        entity_type = "Anything"
+        entity_id = 999
+        multi_entity_update_modes = {"link": "set", "name": "set"}
+        data = {
+            "name": "test",
+            "link": {
+                "name": "test",
+                "url": "http://test.com",
+            },
+        }
+        expected = {
+            "id": 999,
+            "type": "Anything",
+            "fields": [
+                {
+                    "field_name": "name",
+                    "value": "test",
+                    "multi_entity_update_mode": "set",
+                },
+                {
+                    "field_name": "link",
+                    "value": {
+                        "name": "test",
+                        "url": "http://test.com",
+                    },
+                    "multi_entity_update_mode": "set",
+                },
+            ],
+        }
+        sg = api.Shotgun("http://server_path", "script_name", "api_key", connect=False)
+        result = sg._translate_update_params(entity_type, entity_id, data, multi_entity_update_modes)
+        self.assertEqual(result, expected)
+
+    @mock.patch.dict(os.environ, {"SHOTGUN_API_ENABLE_ENTITY_OPTIMIZATION": "1"})
+    def test_related_object_update_optimization_entity(self):
+        entity_type = "Anything"
+        entity_id = 999
+        multi_entity_update_modes = {"project": "set", "link": "set", "name": "set"}
+        data = {
+            "name": "test",
+            "link": {
+                "name": "test",
+                "url": "http://test.com",
+            },
+            "project": {
+                "foo1": "foo1",
+                "bar1": "bar1",
+                "id": 888,
+                "baz1": "baz1",
+                "type": "Project",
+            },
+        }
+        expected = {
+            "id": 999,
+            "type": "Anything",
+            "fields": [
+                {
+                    "field_name": "name",
+                    "value": "test",
+                    "multi_entity_update_mode": "set",
+                },
+                {
+                    "field_name": "link",
+                    "value": {
+                        "name": "test",
+                        "url": "http://test.com",
+                    },
+                    "multi_entity_update_mode": "set",
+                },
+                {
+                    "field_name": "project",
+                    "multi_entity_update_mode": "set",
+                    "value": {
+                        # Entity is optimized with type/id fields.
+                        "id": 888,
+                        "type": "Project",
+                    },
+                },
+            ],
+        }
+        sg = api.Shotgun("http://server_path", "script_name", "api_key", connect=False)
+        result = sg._translate_update_params(entity_type, entity_id, data, multi_entity_update_modes)
+        self.assertEqual(result, expected)
+
+    @mock.patch.dict(os.environ, {"SHOTGUN_API_ENABLE_ENTITY_OPTIMIZATION": "1"})
+    def test_related_object_update_optimization_entity_multi(self):
+        entity_type = "Asset"
+        entity_id = 6626
+        data = {
+            "sg_status_list": "ip",
+            "project": {"id": 70, "type": "Project", "name": "disposable name 70"},
+            "sg_vvv": [
+                {"id": 6441, "type": "Asset", "name": "disposable name 6441"},
+                {"id": 6440, "type": "Asset"},
+            ],
+            "sg_class": {"id": 1, "type": "CustomEntity53", "name": "disposable name 1"},
+        }
+        expected = {
+            "type": "Asset",
+            "id": 6626,
+            "fields": [
+                {"field_name": "sg_status_list", "value": "ip"},
+                {"field_name": "project", "value": {"type": "Project", "id": 70}},
+                {
+                    "field_name": "sg_vvv",
+                    "value": [
+                        {"id": 6441, "type": "Asset"},
+                        {"id": 6440, "type": "Asset"},
+                    ],
+                },
+                {"field_name": "sg_class", "value": {"type": "CustomEntity53", "id": 1}},
+            ],
+        }
+        sg = api.Shotgun("http://server_path", "script_name", "api_key", connect=False)
+        result = sg._translate_update_params(entity_type, entity_id, data, None)
+        self.assertEqual(result, expected)
+
 
 class TestCerts(unittest.TestCase):
     # A dummy bad url provided by Amazon