From 3b528e247989655029135872f62ac6f1396e63fd Mon Sep 17 00:00:00 2001 From: nicholasmhughes Date: Sat, 22 Jan 2022 19:28:03 -0500 Subject: [PATCH 1/7] fixes saltstack/salt#61502 add jinja filters --- changelog/61502.added | 1 + doc/topics/jinja/index.rst | 202 ++++++++++++++++++ salt/utils/data.py | 58 +++++ salt/utils/jinja.py | 62 +++++- .../utils/jinja/test_custom_extensions.py | 135 ++++++++++++ 5 files changed, 457 insertions(+), 1 deletion(-) create mode 100644 changelog/61502.added diff --git a/changelog/61502.added b/changelog/61502.added new file mode 100644 index 000000000000..56fffae9459e --- /dev/null +++ b/changelog/61502.added @@ -0,0 +1 @@ +Add Jinja filters for itertools functions, flatten, and a state template workflow diff --git a/doc/topics/jinja/index.rst b/doc/topics/jinja/index.rst index d66f5c9286a2..6b84685b9f0e 100644 --- a/doc/topics/jinja/index.rst +++ b/doc/topics/jinja/index.rst @@ -279,6 +279,43 @@ that YAML only allows special escapes inside double quotes so use ``yaml_encode`` or ``yaml_dquote``). +.. jinja_ref:: append_dict_key_value + +``dict_to_sls_yaml_params`` +------------------------- + +.. versionadded:: 3005 + +Renders a formatted multi-line YAML string from a Python dictionary. Each +key/value pair in the dictionary will be added as a single-key dictionary +to a list that will then be sent to the YAML formatter. + +Example: + +.. code-block:: jinja + + {% set thing_params = { + "name": "thing", + "changes": True, + "warnings": "OMG! Stuff is happening!" + } + %} + + thing: + test.configurable_test_state: + {{ thing_params | dict_to_sls_yaml_params | indent }} + +Returns: + +.. code-block:: yaml + + thing: + test.configurable_test_state: + - name: thing + - changes: true + - warnings: OMG! Stuff is happening! + + .. jinja_ref:: to_bool ``to_bool`` @@ -651,6 +688,171 @@ Returns: 1, 4 +.. jinja_ref:: flatten + +``flatten`` +------------------------ + +.. versionadded:: 3005 + +Flatten a list. + +.. code-block:: jinja + + {{ [3, [4, 2] ] | flatten }} + # => [3, 4, 2] + +Flatten only the first level of a list: + +.. code-block:: jinja + + {{ [3, [4, [2]] ] | flatten(levels=1) }} + # => [3, 4, [2]] + +Preserve nulls in a list, by default ``flatten`` removes them. + +.. code-block:: jinja + + {{ [3, None, [4, [2]] ] | flatten(levels=1, preserve_nulls=True) }} + # => [3, None, 4, [2]] + + +.. jinja_ref:: combinations + +``combinations`` +------------------------ + +.. versionadded:: 3005 + +Invokes the ``combinations`` function from the ``itertools`` library. + +See the `itertools documentation`_ for more information. + +.. _itertools documentation: https://docs.python.org/3/library/itertools.html#itertools.combinations + +.. code-block:: jinja + + {% for one, two in "ABCD" | combinations(2) %}{{ one~two }} {% endfor %} + # => AB AC AD BC BD CD + + +.. jinja_ref:: combinations_with_replacement + +``combinations_with_replacement`` +------------------------ + +.. versionadded:: 3005 + +Invokes the ``combinations_with_replacement`` function from the ``itertools`` library. + +See the `itertools documentation`_ for more information. + +.. _itertools documentation: https://docs.python.org/3/library/itertools.html#itertools.combinations_with_replacement + +.. code-block:: jinja + + {% for one, two in "ABC" | combinations_with_replacement(2) %}{{ one~two }} {% endfor %} + # => AA AB AC BB BC CC + + +.. jinja_ref:: compress + +``compress`` +------------------------ + +.. versionadded:: 3005 + +Invokes the ``compress`` function from the ``itertools`` library. + +See the `itertools documentation`_ for more information. + +.. _itertools documentation: https://docs.python.org/3/library/itertools.html#itertools.compress + +.. code-block:: jinja + + {% for val in "ABCDEF" | compress([1,0,1,0,1,1]) %}{{ val }} {% endfor %} + # => A C E F + + +.. jinja_ref:: permutations + +``permutations`` +------------------------ + +.. versionadded:: 3005 + +Invokes the ``permutations`` function from the ``itertools`` library. + +See the `itertools documentation`_ for more information. + +.. _itertools documentation: https://docs.python.org/3/library/itertools.html#itertools.permutations + +.. code-block:: jinja + + {% for one, two in "ABCD" | permutations(2) %}{{ one~two }} {% endfor %} + # => AB AC AD BA BC BD CA CB CD DA DB DC + + +.. jinja_ref:: product + +``product`` +------------------------ + +.. versionadded:: 3005 + +Invokes the ``product`` function from the ``itertools`` library. + +See the `itertools documentation`_ for more information. + +.. _itertools documentation: https://docs.python.org/3/library/itertools.html#itertools.product + +.. code-block:: jinja + + {% for one, two in "ABCD" | product("xy") %}{{ one~two }} {% endfor %} + # => Ax Ay Bx By Cx Cy Dx Dy + + +.. jinja_ref:: zip + +``zip`` +------------------------ + +.. versionadded:: 3005 + +Invokes the native Python ``zip`` function. + +The ``zip`` function returns a zip object, which is an iterator of tuples where +the first item in each passed iterator is paired together, and then the second +item in each passed iterator are paired together etc. + +If the passed iterators have different lengths, the iterator with the least +items decides the length of the new iterator. + +.. code-block:: jinja + + {% for one, two in "ABCD" | zip("xy") %}{{ one~two }} {% endfor %} + # => Ax By + + +.. jinja_ref:: zip_longest + +``zip_longest`` +------------------------ + +.. versionadded:: 3005 + +Invokes the ``zip_longest`` function from the ``itertools`` library. + +See the `itertools documentation`_ for more information. + +.. _itertools documentation: https://docs.python.org/3/library/itertools.html#itertools.zip_longest + +.. code-block:: jinja + + {% for one, two in "ABCD" | zip_longest("xy", fillvalue="-") %}{{ one~two }} {% endfor %} + # => Ax By C- D- + + .. jinja_ref:: method_call ``method_call`` diff --git a/salt/utils/data.py b/salt/utils/data.py index 8559c5bf3c7d..019d552cb67d 100644 --- a/salt/utils/data.py +++ b/salt/utils/data.py @@ -1541,3 +1541,61 @@ def _append_placeholder(value_dict, key): else: return [{"value": default if obj is not None else obj}] return res + + +@jinja_filter("flatten") +def flatten(data, levels=None, preserve_nulls=False): + """ + .. versionadded:: 3005 + + Flatten a list. + + :param data: A list to flatten + + :param levels: The number of levels in sub-lists to descend + + :param preserve_nulls: Preserve nulls in a list, by default flatten removes them + + :returns: A flat(ter) list of values + + .. code-block:: jinja + + {{ [3, [4, 2] ] | flatten }} + # => [3, 4, 2] + + Flatten only the first level of a list: + + .. code-block:: jinja + + {{ [3, [4, [2]] ] | flatten(levels=1) }} + # => [3, 4, [2]] + + Preserve nulls in a list, by default flatten removes them. + + .. code-block:: jinja + + {{ [3, None, [4, [2]] ] | flatten(levels=1, preserve_nulls=True) }} + # => [3, None, 4, [2]] + """ + ret = [] + + for element in data: + if not preserve_nulls and element in (None, "None", "null"): + # ignore null items + continue + elif is_iter(element): + if levels is None: + ret.extend(flatten(element, preserve_nulls=preserve_nulls)) + elif levels >= 1: + # decrement as we go down the stack + ret.extend( + flatten( + element, levels=(int(levels) - 1), preserve_nulls=preserve_nulls + ) + ) + else: + ret.append(element) + else: + ret.append(element) + + return ret diff --git a/salt/utils/jinja.py b/salt/utils/jinja.py index 0cb70bf64aec..4c430b5ccf32 100644 --- a/salt/utils/jinja.py +++ b/salt/utils/jinja.py @@ -4,6 +4,7 @@ import atexit +import itertools import logging import os.path import pipes @@ -877,6 +878,39 @@ class SerializerExtension(Extension): unique = ['foo', 'bar'] + ** Salt State Parameter Format Filters ** + + .. versionadded:: 3005 + + Renders a formatted multi-line YAML string from a Python dictionary. Each + key/value pair in the dictionary will be added as a single-key dictionary + to a list that will then be sent to the YAML formatter. + + For example: + + .. code-block:: jinja + + {% set thing_params = { + "name": "thing", + "changes": True, + "warnings": "OMG! Stuff is happening!" + } + %} + + thing: + test.configurable_test_state: + {{ thing_params | dict_to_sls_yaml_params | indent }} + + will be rendered as:: + + .. code-block:: yaml + + thing: + test.configurable_test_state: + - name: thing + - changes: true + - warnings: OMG! Stuff is happening! + .. _`import tag`: https://jinja.palletsprojects.com/en/2.11.x/templates/#import ''' @@ -901,6 +935,14 @@ def __init__(self, environment): "load_yaml": self.load_yaml, "load_json": self.load_json, "load_text": self.load_text, + "dict_to_sls_yaml_params": self.dict_to_sls_yaml_params, + "combinations": itertools.combinations, + "combinations_with_replacement": itertools.combinations_with_replacement, + "compress": itertools.compress, + "permutations": itertools.permutations, + "product": itertools.product, + "zip": zip, + "zip_longest": itertools.zip_longest, } ) @@ -1169,4 +1211,22 @@ def parse_import(self, parser, converter): parser, import_node.template, "import_{}".format(converter), body, lineno ) - # pylint: enable=E1120,E1121 + def dict_to_sls_yaml_params(self, value, flow_style=False): + """ + .. versionadded:: 3005 + + Render a formatted multi-line YAML string from a Python dictionary. Each + key/value pair in the dictionary will be added as a single-key dictionary + to a list that will then be sent to the YAML formatter. + + :param value: Python dictionary representing Salt state parameters + + :param flow_style: Setting flow_style to False will enforce indentation + mode + + :returns: Formatted SLS YAML string rendered with newlines and + indentation + """ + return self.format_yaml( + [{key: val} for key, val in value.items()], flow_style=flow_style + ) diff --git a/tests/pytests/unit/utils/jinja/test_custom_extensions.py b/tests/pytests/unit/utils/jinja/test_custom_extensions.py index c1c4451d83ef..58a2f4af09b1 100644 --- a/tests/pytests/unit/utils/jinja/test_custom_extensions.py +++ b/tests/pytests/unit/utils/jinja/test_custom_extensions.py @@ -1070,3 +1070,138 @@ def test_json_query(minion_opts, local_salt): dict(opts=minion_opts, saltenv="test", salt=local_salt), ) assert rendered == "2" + + +def test_flatten_simple(minion_opts, local_salt): + """ + Test the `flatten` Jinja filter. + """ + rendered = render_jinja_tmpl( + "{{ [1, 2, [3]] | flatten }}", + dict(opts=minion_opts, saltenv="test", salt=local_salt), + ) + assert rendered == "[1, 2, 3]" + + +def test_flatten_single_level(minion_opts, local_salt): + """ + Test the `flatten` Jinja filter. + """ + rendered = render_jinja_tmpl( + "{{ [1, 2, [None, 3, [4]]] | flatten(levels=1) }}", + dict(opts=minion_opts, saltenv="test", salt=local_salt), + ) + assert rendered == "[1, 2, 3, [4]]" + + +def test_flatten_preserve_nulls(minion_opts, local_salt): + """ + Test the `flatten` Jinja filter. + """ + rendered = render_jinja_tmpl( + "{{ [1, 2, [None, 3, [4]]] | flatten(preserve_nulls=True) }}", + dict(opts=minion_opts, saltenv="test", salt=local_salt), + ) + assert rendered == "[1, 2, None, 3, 4]" + + +def test_dict_to_sls_yaml_params(minion_opts, local_salt): + """ + Test the `dict_to_sls_yaml_params` Jinja filter. + """ + expected = ( + "- name: donkey\n" + + "- list:\n" + + " - one\n" + + " - two\n" + + "- dict:\n" + + " one: two\n" + + "- nested:\n" + + " - one\n" + + " - two: three" + ) + source = ( + "{% set myparams = {'name': 'donkey', 'list': ['one', 'two'], 'dict': {'one': 'two'}, 'nested': ['one', {'two': 'three'}]} %}" + + "{{ myparams | dict_to_sls_yaml_params }}" + ) + rendered = render_jinja_tmpl( + source, dict(opts=minion_opts, saltenv="test", salt=local_salt) + ) + assert rendered == expected + + +def test_combinations(minion_opts, local_salt): + """ + Test the `combinations` Jinja filter. + """ + rendered = render_jinja_tmpl( + "{% for one, two in 'ABCD' | combinations(2) %}{{ one~two }} {% endfor %}", + dict(opts=minion_opts, saltenv="test", salt=local_salt), + ) + assert rendered == "AB AC AD BC BD CD " + + +def test_combinations_with_replacement(minion_opts, local_salt): + """ + Test the `combinations_with_replacement` Jinja filter. + """ + rendered = render_jinja_tmpl( + "{% for one, two in 'ABC' | combinations_with_replacement(2) %}{{ one~two }} {% endfor %}", + dict(opts=minion_opts, saltenv="test", salt=local_salt), + ) + assert rendered == "AA AB AC BB BC CC " + + +def test_compress(minion_opts, local_salt): + """ + Test the `compress` Jinja filter. + """ + rendered = render_jinja_tmpl( + "{% for val in 'ABCDEF' | compress([1,0,1,0,1,1]) %}{{ val }} {% endfor %}", + dict(opts=minion_opts, saltenv="test", salt=local_salt), + ) + assert rendered == "A C E F " + + +def test_permutations(minion_opts, local_salt): + """ + Test the `permutations` Jinja filter. + """ + rendered = render_jinja_tmpl( + "{% for one, two in 'ABCD' | permutations(2) %}{{ one~two }} {% endfor %}", + dict(opts=minion_opts, saltenv="test", salt=local_salt), + ) + assert rendered == "AB AC AD BA BC BD CA CB CD DA DB DC " + + +def test_product(minion_opts, local_salt): + """ + Test the `product` Jinja filter. + """ + rendered = render_jinja_tmpl( + "{% for one, two in 'ABCD' | product('xy') %}{{ one~two }} {% endfor %}", + dict(opts=minion_opts, saltenv="test", salt=local_salt), + ) + assert rendered == "Ax Ay Bx By Cx Cy Dx Dy " + + +def test_zip(minion_opts, local_salt): + """ + Test the `zip` Jinja filter. + """ + rendered = render_jinja_tmpl( + "{% for one, two in 'ABCD' | zip('xy') %}{{ one~two }} {% endfor %}", + dict(opts=minion_opts, saltenv="test", salt=local_salt), + ) + assert rendered == "Ax By " + + +def test_zip_longest(minion_opts, local_salt): + """ + Test the `zip_longest` Jinja filter. + """ + rendered = render_jinja_tmpl( + "{% for one, two in 'ABCD' | zip_longest('xy', fillvalue='-') %}{{ one~two }} {% endfor %}", + dict(opts=minion_opts, saltenv="test", salt=local_salt), + ) + assert rendered == "Ax By C- D- " From 27aa052b9a35cccb67410503822e05f014388f82 Mon Sep 17 00:00:00 2001 From: nicholasmhughes Date: Sat, 22 Jan 2022 20:09:11 -0500 Subject: [PATCH 2/7] title underline too short --- doc/topics/jinja/index.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/topics/jinja/index.rst b/doc/topics/jinja/index.rst index 6b84685b9f0e..9b88897d3a09 100644 --- a/doc/topics/jinja/index.rst +++ b/doc/topics/jinja/index.rst @@ -282,7 +282,7 @@ use ``yaml_encode`` or ``yaml_dquote``). .. jinja_ref:: append_dict_key_value ``dict_to_sls_yaml_params`` -------------------------- +--------------------------- .. versionadded:: 3005 From edde840f1ecbbea3d1dbc631b2fc332e918fe08c Mon Sep 17 00:00:00 2001 From: nicholasmhughes Date: Sat, 22 Jan 2022 20:10:40 -0500 Subject: [PATCH 3/7] copy paste error --- doc/topics/jinja/index.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/topics/jinja/index.rst b/doc/topics/jinja/index.rst index 9b88897d3a09..da4e4c8a6a3f 100644 --- a/doc/topics/jinja/index.rst +++ b/doc/topics/jinja/index.rst @@ -279,7 +279,7 @@ that YAML only allows special escapes inside double quotes so use ``yaml_encode`` or ``yaml_dquote``). -.. jinja_ref:: append_dict_key_value +.. jinja_ref:: dict_to_sls_yaml_params ``dict_to_sls_yaml_params`` --------------------------- From 51e711ef6cb08010b9357e25e2ed6a0f83487af6 Mon Sep 17 00:00:00 2001 From: nicholasmhughes Date: Sat, 22 Jan 2022 20:22:41 -0500 Subject: [PATCH 4/7] fixing doc underlines --- doc/topics/jinja/index.rst | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/doc/topics/jinja/index.rst b/doc/topics/jinja/index.rst index da4e4c8a6a3f..ba136e26791b 100644 --- a/doc/topics/jinja/index.rst +++ b/doc/topics/jinja/index.rst @@ -691,7 +691,7 @@ Returns: .. jinja_ref:: flatten ``flatten`` ------------------------- +----------- .. versionadded:: 3005 @@ -720,7 +720,7 @@ Preserve nulls in a list, by default ``flatten`` removes them. .. jinja_ref:: combinations ``combinations`` ------------------------- +---------------- .. versionadded:: 3005 @@ -739,7 +739,7 @@ See the `itertools documentation`_ for more information. .. jinja_ref:: combinations_with_replacement ``combinations_with_replacement`` ------------------------- +--------------------------------- .. versionadded:: 3005 @@ -758,7 +758,7 @@ See the `itertools documentation`_ for more information. .. jinja_ref:: compress ``compress`` ------------------------- +------------ .. versionadded:: 3005 @@ -777,7 +777,7 @@ See the `itertools documentation`_ for more information. .. jinja_ref:: permutations ``permutations`` ------------------------- +---------------- .. versionadded:: 3005 @@ -796,7 +796,7 @@ See the `itertools documentation`_ for more information. .. jinja_ref:: product ``product`` ------------------------- +----------- .. versionadded:: 3005 @@ -815,7 +815,7 @@ See the `itertools documentation`_ for more information. .. jinja_ref:: zip ``zip`` ------------------------- +------- .. versionadded:: 3005 @@ -837,7 +837,7 @@ items decides the length of the new iterator. .. jinja_ref:: zip_longest ``zip_longest`` ------------------------- +--------------- .. versionadded:: 3005 From dbde20f3811f1b98b9ed6ec7fb0326f11baa2057 Mon Sep 17 00:00:00 2001 From: nicholasmhughes Date: Sat, 22 Jan 2022 21:09:55 -0500 Subject: [PATCH 5/7] fix itertools doc links --- doc/topics/jinja/index.rst | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/doc/topics/jinja/index.rst b/doc/topics/jinja/index.rst index ba136e26791b..54b447d63514 100644 --- a/doc/topics/jinja/index.rst +++ b/doc/topics/jinja/index.rst @@ -728,8 +728,6 @@ Invokes the ``combinations`` function from the ``itertools`` library. See the `itertools documentation`_ for more information. -.. _itertools documentation: https://docs.python.org/3/library/itertools.html#itertools.combinations - .. code-block:: jinja {% for one, two in "ABCD" | combinations(2) %}{{ one~two }} {% endfor %} @@ -747,8 +745,6 @@ Invokes the ``combinations_with_replacement`` function from the ``itertools`` li See the `itertools documentation`_ for more information. -.. _itertools documentation: https://docs.python.org/3/library/itertools.html#itertools.combinations_with_replacement - .. code-block:: jinja {% for one, two in "ABC" | combinations_with_replacement(2) %}{{ one~two }} {% endfor %} @@ -766,8 +762,6 @@ Invokes the ``compress`` function from the ``itertools`` library. See the `itertools documentation`_ for more information. -.. _itertools documentation: https://docs.python.org/3/library/itertools.html#itertools.compress - .. code-block:: jinja {% for val in "ABCDEF" | compress([1,0,1,0,1,1]) %}{{ val }} {% endfor %} @@ -785,8 +779,6 @@ Invokes the ``permutations`` function from the ``itertools`` library. See the `itertools documentation`_ for more information. -.. _itertools documentation: https://docs.python.org/3/library/itertools.html#itertools.permutations - .. code-block:: jinja {% for one, two in "ABCD" | permutations(2) %}{{ one~two }} {% endfor %} @@ -804,8 +796,6 @@ Invokes the ``product`` function from the ``itertools`` library. See the `itertools documentation`_ for more information. -.. _itertools documentation: https://docs.python.org/3/library/itertools.html#itertools.product - .. code-block:: jinja {% for one, two in "ABCD" | product("xy") %}{{ one~two }} {% endfor %} From 422e499349a1ac4ffbb9198fe11d10afb695c8cc Mon Sep 17 00:00:00 2001 From: nicholasmhughes Date: Sun, 23 Jan 2022 10:04:53 -0500 Subject: [PATCH 6/7] dicts are unordered so tests need to check permutations --- .../utils/jinja/test_custom_extensions.py | 20 ++++++++----------- 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/tests/pytests/unit/utils/jinja/test_custom_extensions.py b/tests/pytests/unit/utils/jinja/test_custom_extensions.py index 58a2f4af09b1..30d210c3c483 100644 --- a/tests/pytests/unit/utils/jinja/test_custom_extensions.py +++ b/tests/pytests/unit/utils/jinja/test_custom_extensions.py @@ -3,6 +3,7 @@ """ import ast +import itertools import os import pprint import random @@ -1109,17 +1110,12 @@ def test_dict_to_sls_yaml_params(minion_opts, local_salt): """ Test the `dict_to_sls_yaml_params` Jinja filter. """ - expected = ( - "- name: donkey\n" - + "- list:\n" - + " - one\n" - + " - two\n" - + "- dict:\n" - + " one: two\n" - + "- nested:\n" - + " - one\n" - + " - two: three" - ) + expected = [ + "- name: donkey", + "- list:\n - one\n - two", + "- dict:\n one: two", + "- nested:\n - one\n - two: three", + ] source = ( "{% set myparams = {'name': 'donkey', 'list': ['one', 'two'], 'dict': {'one': 'two'}, 'nested': ['one', {'two': 'three'}]} %}" + "{{ myparams | dict_to_sls_yaml_params }}" @@ -1127,7 +1123,7 @@ def test_dict_to_sls_yaml_params(minion_opts, local_salt): rendered = render_jinja_tmpl( source, dict(opts=minion_opts, saltenv="test", salt=local_salt) ) - assert rendered == expected + assert rendered in ["\n".join(combo) for combo in itertools.permutations(expected)] def test_combinations(minion_opts, local_salt): From 8c5368bd1102be7783f30786147ea2358171853c Mon Sep 17 00:00:00 2001 From: nicholasmhughes Date: Thu, 27 Jan 2022 09:58:49 -0500 Subject: [PATCH 7/7] detect reference cycles in flatten function --- salt/utils/data.py | 21 +++++++++++++++++---- tests/pytests/unit/utils/test_data.py | 11 +++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/salt/utils/data.py b/salt/utils/data.py index 019d552cb67d..0b3634d84240 100644 --- a/salt/utils/data.py +++ b/salt/utils/data.py @@ -1544,7 +1544,7 @@ def _append_placeholder(value_dict, key): @jinja_filter("flatten") -def flatten(data, levels=None, preserve_nulls=False): +def flatten(data, levels=None, preserve_nulls=False, _ids=None): """ .. versionadded:: 3005 @@ -1554,7 +1554,11 @@ def flatten(data, levels=None, preserve_nulls=False): :param levels: The number of levels in sub-lists to descend - :param preserve_nulls: Preserve nulls in a list, by default flatten removes them + :param preserve_nulls: Preserve nulls in a list, by default flatten removes + them + + :param _ids: Parameter used internally within the function to detect + reference cycles. :returns: A flat(ter) list of values @@ -1577,6 +1581,12 @@ def flatten(data, levels=None, preserve_nulls=False): {{ [3, None, [4, [2]] ] | flatten(levels=1, preserve_nulls=True) }} # => [3, None, 4, [2]] """ + if _ids is None: + _ids = set() + if id(data) in _ids: + raise RecursionError("Reference cycle detected. Check input list.") + _ids.add(id(data)) + ret = [] for element in data: @@ -1585,12 +1595,15 @@ def flatten(data, levels=None, preserve_nulls=False): continue elif is_iter(element): if levels is None: - ret.extend(flatten(element, preserve_nulls=preserve_nulls)) + ret.extend(flatten(element, preserve_nulls=preserve_nulls, _ids=_ids)) elif levels >= 1: # decrement as we go down the stack ret.extend( flatten( - element, levels=(int(levels) - 1), preserve_nulls=preserve_nulls + element, + levels=(int(levels) - 1), + preserve_nulls=preserve_nulls, + _ids=_ids, ) ) else: diff --git a/tests/pytests/unit/utils/test_data.py b/tests/pytests/unit/utils/test_data.py index bd8989621a4b..97fe8b9de0b7 100644 --- a/tests/pytests/unit/utils/test_data.py +++ b/tests/pytests/unit/utils/test_data.py @@ -70,3 +70,14 @@ def test_get_value_simple_type_path(): def test_get_value_None_path(): assert [{"value": None}] == salt.utils.data.get_value({"a": None}, "a:b", []) + + +def test_flatten_recursion_error(): + """ + Test the flatten function for reference cycle detection + """ + data = [1, 2, 3, [4]] + data.append(data) + with pytest.raises(RecursionError) as err: + salt.utils.data.flatten(data) + assert str(err.value) == "Reference cycle detected. Check input list."