Fix bug in prune_dict where empty dict and list would be removed even in strict mode (#32573)

vandonr-amz · web-flow · commit bbd290266c0a · 2023-07-19T15:43:21.000-04:00
diff --git a/airflow/providers/amazon/aws/utils/__init__.py b/airflow/providers/amazon/aws/utils/__init__.py
@@ -21,13 +21,28 @@
 from datetime import datetime
 from enum import Enum
 
+from airflow.utils.helpers import prune_dict
 from airflow.version import version
 
 log = logging.getLogger(__name__)
 
 
 def trim_none_values(obj: dict):
-    return {key: val for key, val in obj.items() if val is not None}
+    from packaging.version import Version
+
+    from airflow.version import version
+
+    if Version(version) < Version("2.7"):
+        # before version 2.7, the behavior is not the same.
+        # Empty dict and lists are removed from the given dict.
+        return {key: val for key, val in obj.items() if val is not None}
+    else:
+        # once airflow 2.6 rolls out of compatibility support for provider packages,
+        # we can replace usages of this method with the core one in our code,
+        # and uncomment this warning for users who may use it.
+        # warnings.warn("use airflow.utils.helpers.prune_dict() instead",
+        #     AirflowProviderDeprecationWarning, stacklevel=2)
+        return prune_dict(obj)
 
 
 def datetime_to_epoch(date_time: datetime) -> int:
diff --git a/airflow/utils/helpers.py b/airflow/utils/helpers.py
@@ -341,7 +341,7 @@ def is_empty(x):
                 continue
             elif isinstance(v, (list, dict)):
                 new_val = prune_dict(v, mode=mode)
-                if new_val:
+                if not is_empty(new_val):
                     new_dict[k] = new_val
             else:
                 new_dict[k] = v
@@ -353,7 +353,7 @@ def is_empty(x):
                 continue
             elif isinstance(v, (list, dict)):
                 new_val = prune_dict(v, mode=mode)
-                if new_val:
+                if not is_empty(new_val):
                     new_list.append(new_val)
             else:
                 new_list.append(v)
diff --git a/tests/providers/amazon/aws/utils/test_utils.py b/tests/providers/amazon/aws/utils/test_utils.py
@@ -26,7 +26,6 @@
     datetime_to_epoch_ms,
     datetime_to_epoch_us,
     get_airflow_version,
-    trim_none_values,
 )
 
 DT = datetime(2000, 1, 1, tzinfo=pytz.UTC)
@@ -37,17 +36,6 @@ class EnumTest(_StringCompareEnum):
     FOO = "FOO"
 
 
-def test_trim_none_values():
-    input_object = {
-        "test": "test",
-        "empty": None,
-    }
-    expected_output_object = {
-        "test": "test",
-    }
-    assert trim_none_values(input_object) == expected_output_object
-
-
 def test_datetime_to_epoch():
     assert datetime_to_epoch(DT) == EPOCH
 
diff --git a/tests/utils/test_helpers.py b/tests/utils/test_helpers.py
@@ -309,6 +309,8 @@ def assert_at_most_one(true=0, truthy=0, false=0, falsy=0, notset=0):
                     "c": {"b": "", "c": "hi", "d": ["", 0, "1"]},
                     "d": ["", 0, "1"],
                     "e": ["", 0, {"b": "", "c": "hi", "d": ["", 0, "1"]}, ["", 0, "1"], [""]],
+                    "f": {},
+                    "g": [""],
                 },
             ),
             (
@@ -324,7 +326,7 @@ def assert_at_most_one(true=0, truthy=0, false=0, falsy=0, notset=0):
     def test_prune_dict(self, mode, expected):
         l1 = ["", 0, "1", None]
         d1 = {"a": None, "b": "", "c": "hi", "d": l1}
-        d2 = {"a": None, "b": "", "c": d1, "d": l1, "e": [None, "", 0, d1, l1, [""]]}
+        d2 = {"a": None, "b": "", "c": d1, "d": l1, "e": [None, "", 0, d1, l1, [""]], "f": {}, "g": [""]}
         assert prune_dict(d2, mode=mode) == expected
 
 
