SG-39225 Support specific keys for entity optimization (#423)

* Support specific keys for entity optimization

* Update conditions

* Update test

* Support lists without dicts on optimization

* Improved readability

* Improvements by feedback

* More granular optimization

* Use simple conditions instead of set intersection

* Include id and type checks

* Make recursive and add more test cases

Author: carlos-villavicencio-adsk

shotgun_api3/shotgun.py

@@ -1194,7 +1194,7 @@ def _translate_update_params(
         def optimize_field(field_dict):
             if SHOTGUN_API_DISABLE_ENTITY_OPTIMIZATION:
                 return field_dict
-            return {k: _get_type_and_id_from_value(v) for k, v in field_dict.items()}
+            return {k: _optimize_filter_field(v) for k, v in field_dict.items()}
 
         full_fields = self._dict_to_list(
             data,
@@ -4799,7 +4799,7 @@ def _translate_filters_simple(sg_filter):
         and condition["relation"] in ["is", "is_not", "in", "not_in"]
         and isinstance(values[0], dict)
     ):
-        values = [_get_type_and_id_from_value(v) for v in values]
+        values = [_optimize_filter_field(v) for v in values]
 
     condition["values"] = values
 
@@ -4813,17 +4813,30 @@ def _version_str(version) -> str:
     return ".".join(map(str, version))
 
 
-def _get_type_and_id_from_value(value):
+def _optimize_filter_field(
+    field_value: Union[dict, list], recursive: bool = True
+) -> Union[dict, list]:
     """
-    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.
+    For an FPT entity, returns a new dictionary with only the type,
+    id, and other allowed keys.
+    If case of any processing error, the original dictionary is returned.
+
+    At least `type` and `id` keys are required to do the optimization
     """
-    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
+    allowed_keys = {
+        "id",
+        "type",
+        "url",
+        "name",
+        "content_type",
+        "local_path",
+        "storage",
+        "relative_path",
+    }
+    if isinstance(field_value, dict) and "id" in field_value and "type" in field_value:
+        return {key: field_value[key] for key in allowed_keys if key in field_value}
+
+    elif recursive and isinstance(field_value, list):
+        return [_optimize_filter_field(fv, recursive=False) for fv in field_value]
+
+    return field_value

tests/test_unit.py

@@ -673,33 +673,88 @@ def test_related_object_update_optimization_entity_multi(self):
         entity_id = 6626
         data = {
             "sg_status_list": "ip",
-            "project": {"id": 70, "type": "Project", "name": "disposable name 70"},
+            "project": {"id": 70, "type": "Project", "name": "important name 70"},
             "sg_vvv": [
-                {"id": 6441, "type": "Asset", "name": "disposable name 6441"},
                 {"id": 6440, "type": "Asset"},
+                {"id": 6441, "type": "Asset", "custom_name": "disposable name 6441"},
+                {
+                    # To be kept
+                    "id": 6442,
+                    "type": "Asset",
+                    "url": "http://test.com/asset/6442",
+                    # to be removed
+                    "custom_name": "disposable name 1",
+                    "custom_name2": "disposable name 2",
+                    "custom_name3": "disposable name 3",
+                    "custom_name4": "disposable name 4",
+                },
+                {
+                    "sg_nested": {
+                        "level1": {
+                            "level2": {"id": 123, "type": "Entity", "foo": "bar"}
+                        }
+                    }
+                },
             ],
             "sg_class": {
+                # To be kept
                 "id": 1,
                 "type": "CustomEntity53",
-                "name": "disposable name 1",
+                "url": "http://test.com",
+                "name": "important class name",
+                "local_path": "/some/local/path",
+                # to be removed
+                "custom_name": "disposable name 1",
+                "custom_name2": "disposable name 2",
+                "custom_name3": "disposable name 3",
+                "custom_name4": "disposable name 4",
             },
         }
         expected = {
             "type": "Asset",
             "id": 6626,
             "fields": [
                 {"field_name": "sg_status_list", "value": "ip"},
-                {"field_name": "project", "value": {"type": "Project", "id": 70}},
+                {
+                    "field_name": "project",
+                    "value": {
+                        "type": "Project",
+                        "id": 70,
+                        "name": "important name 70",
+                    },
+                },
                 {
                     "field_name": "sg_vvv",
                     "value": [
-                        {"id": 6441, "type": "Asset"},
                         {"id": 6440, "type": "Asset"},
+                        {"id": 6441, "type": "Asset"},
+                        {
+                            "id": 6442,
+                            "type": "Asset",
+                            "url": "http://test.com/asset/6442",
+                        },
+                        {
+                            "sg_nested": {
+                                "level1": {
+                                    "level2": {
+                                        "id": 123,
+                                        "type": "Entity",
+                                        "foo": "bar",
+                                    }
+                                }
+                            }
+                        },
                     ],
                 },
                 {
                     "field_name": "sg_class",
-                    "value": {"type": "CustomEntity53", "id": 1},
+                    "value": {
+                        "type": "CustomEntity53",
+                        "id": 1,
+                        "name": "important class name",
+                        "url": "http://test.com",
+                        "local_path": "/some/local/path",
+                    },
                 },
             ],
         }