alter flatten function to work on dict only. not a list of dicts

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

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",
     }