From 1938c316216dd45827e70c67ae66315bcf7d11cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felipe=20Sodr=C3=A9?= <felipe.b4rros@gmail.com>
Date: Tue, 13 Feb 2024 17:28:30 -0300
Subject: [PATCH] alter flatten function to work on dict only. not a list of
 dicts

---
 crossfire/__init__.py | 13 +++++++------
 tests/test_flatten.py | 12 ++++++------
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/crossfire/__init__.py b/crossfire/__init__.py
index 17d7d43..aae3f53 100644
--- a/crossfire/__init__.py
+++ b/crossfire/__init__.py
@@ -48,9 +48,10 @@ def flatten(data, nested_columns=None):
     nested_columns = set(nested_columns or NESTED_COLUMNS)
     if not nested_columns.issubset(NESTED_COLUMNS):
         raise NestedColumnError(nested_columns)
-    if isinstance(data, list):
-        keys = set(data[0].keys()) & nested_columns
-        for item in data:
-            for key in keys:
-                for k, v in item.get(key).items():
-                    return f"{key}_{k}", v
+    if isinstance(data, dict):
+        keys = set(data.keys()) & nested_columns
+        flattened_dict = {}
+        for key in keys:
+            for k, v in data.get(key).items():
+                flattened_dict.update({f"{key}_{k}": v})
+            return flattened_dict
diff --git a/tests/test_flatten.py b/tests/test_flatten.py
index f7b0ef7..7c47f4b 100644
--- a/tests/test_flatten.py
+++ b/tests/test_flatten.py
@@ -5,12 +5,13 @@
 
 from crossfire import NestedColumnError, flatten
 
-DICT_DATA = [
-    {"answer": 42, "contextInfo": {"context1": "info1", "context2": "info2"}}
-]
-PD_DATA = DataFrame(DICT_DATA)
+DICT_DATA = {
+    "answer": 42,
+    "contextInfo": {"context1": "info1", "context2": "info2"},
+}
+PD_DATA = DataFrame([DICT_DATA])
 GEOMETRY = [Point(4, 2)]
-GEOPD_DATA = GeoDataFrame(DICT_DATA, crs="EPSG:4326", geometry=GEOMETRY)
+GEOPD_DATA = GeoDataFrame([DICT_DATA], crs="EPSG:4326", geometry=GEOMETRY)
 
 
 def teste_flatten_wrong_nested_columns_value_error():
@@ -24,7 +25,6 @@ def test_flatten_dict():
         DICT_DATA, nested_columns=["contextInfo", "transports"]
     )
     assert flattened_dict == {
-        "answer": 42,
         "contextInfo_context1": "info1",
         "contextInfo_context2": "info2",
     }