Skip to content

Commit

Permalink
Fix bug in prune_dict where empty dict and list would be removed even…
Browse files Browse the repository at this point in the history
… in strict mode (#32573)
  • Loading branch information
vandonr-amz committed Jul 19, 2023
1 parent 94122d1 commit bbd2902
Show file tree
Hide file tree
Showing 4 changed files with 21 additions and 16 deletions.
17 changes: 16 additions & 1 deletion airflow/providers/amazon/aws/utils/__init__.py
Expand Up @@ -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:
Expand Down
4 changes: 2 additions & 2 deletions airflow/utils/helpers.py
Expand Up @@ -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
Expand All @@ -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)
Expand Down
12 changes: 0 additions & 12 deletions tests/providers/amazon/aws/utils/test_utils.py
Expand Up @@ -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)
Expand All @@ -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

Expand Down
4 changes: 3 additions & 1 deletion tests/utils/test_helpers.py
Expand Up @@ -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": [""],
},
),
(
Expand All @@ -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


Expand Down

0 comments on commit bbd2902

Please sign in to comment.