From a19c27426c7008769530d72420a88a68b9aed548 Mon Sep 17 00:00:00 2001
From: Jon Mease <jon@plotly.com>
Date: Fri, 6 Aug 2021 15:34:57 -0400
Subject: [PATCH] Preserve schema key order when flattening a dict grouping
 (#1701)

* Preserve schema key order when flattening a dict grouping
* Python 3.6 compatibility
---
 dash/_grouping.py                |  8 +++-----
 dash/dash.py                     |  5 +----
 tests/unit/dash/test_grouping.py | 11 +++++++++++
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/dash/_grouping.py b/dash/_grouping.py
index cb71cd8e1c..89d62de8aa 100644
--- a/dash/_grouping.py
+++ b/dash/_grouping.py
@@ -30,6 +30,8 @@ def flatten_grouping(grouping, schema=None):
     """
     if schema is None:
         schema = grouping
+    else:
+        validate_grouping(grouping, schema)
 
     if isinstance(schema, (tuple, list)):
         return [
@@ -39,11 +41,7 @@ def flatten_grouping(grouping, schema=None):
         ]
 
     if isinstance(schema, dict):
-        return [
-            g
-            for group_el, schema_el in zip(grouping.values(), schema.values())
-            for g in flatten_grouping(group_el, schema_el)
-        ]
+        return [g for k in schema for g in flatten_grouping(grouping[k], schema[k])]
 
     return [grouping]
 
diff --git a/dash/dash.py b/dash/dash.py
index c7fbe64d79..73b150ee90 100644
--- a/dash/dash.py
+++ b/dash/dash.py
@@ -55,7 +55,6 @@
 from . import _watch
 from ._grouping import (
     flatten_grouping,
-    validate_grouping,
     map_grouping,
     make_grouping_by_index,
     grouping_len,
@@ -1075,11 +1074,9 @@ def add_context(*args, **kwargs):
                         # list or tuple
                         output_value = list(output_value)
 
+                    # Flatten grouping and validate grouping structure
                     flat_output_values = flatten_grouping(output_value, output)
 
-                # Validate grouping structure
-                validate_grouping(output_value, output)
-
                 _validate.validate_multi_return(
                     output_spec, flat_output_values, callback_id
                 )
diff --git a/tests/unit/dash/test_grouping.py b/tests/unit/dash/test_grouping.py
index 215afcc7aa..d773f0596e 100644
--- a/tests/unit/dash/test_grouping.py
+++ b/tests/unit/dash/test_grouping.py
@@ -45,6 +45,17 @@ def test_flatten_dict(dict_grouping_size):
     assert len(result) == grouping_len(grouping)
 
 
+def test_flatten_dict_key_order(dict_grouping_size):
+    grouping, size = dict_grouping_size
+    expected = list(range(size))
+
+    # Reverse key order of value dict to make sure order is preserved
+    rev_grouping = {k: grouping[k] for k in reversed(list(grouping.keys()))}
+    result = flatten_grouping(rev_grouping, grouping)
+    assert expected == result
+    assert len(result) == grouping_len(grouping)
+
+
 def test_flatten_mixed(mixed_grouping_size):
     grouping, size = mixed_grouping_size
     expected = list(range(size))