minor change to FilePrediction.to_json structure

src/stratigraphy/util/predictions.py

@@ -209,20 +209,16 @@ def to_json(self) -> dict:
             dict: The object as a dictionary.
         """
         return {
-            self.file_name: {
-                "metadata": self.metadata.to_json(),
-                "layers": [layer.to_json() for layer in self.layers] if self.layers is not None else [],
-                "depths_materials_column_pairs": [
-                    dmc_pair.to_json() for dmc_pair in self.depths_materials_columns_pairs
-                ]
-                if self.depths_materials_columns_pairs is not None
-                else [],
-                "page_dimensions": self.metadata.page_dimensions,  # TODO: Remove, already in metadata
-                "groundwater": [entry.to_json() for entry in self.groundwater_entries]
-                if self.groundwater_entries is not None
-                else [],
-                "file_name": self.file_name,
-            }
+            "metadata": self.metadata.to_json(),
+            "layers": [layer.to_json() for layer in self.layers] if self.layers is not None else [],
+            "depths_materials_column_pairs": [dmc_pair.to_json() for dmc_pair in self.depths_materials_columns_pairs]
+            if self.depths_materials_columns_pairs is not None
+            else [],
+            "page_dimensions": self.metadata.page_dimensions,  # TODO: Remove, already in metadata
+            "groundwater": [entry.to_json() for entry in self.groundwater_entries]
+            if self.groundwater_entries is not None
+            else [],
+            "file_name": self.file_name,
         }
 
 
@@ -254,7 +250,7 @@ def to_json(self) -> dict:
         Returns:
             dict: A dictionary representation of the object.
         """
-        return {k: v for fp in self.file_predictions_list for k, v in fp.to_json().items()}
+        return {fp.file_name: fp.to_json() for fp in self.file_predictions_list}
 
     @classmethod
     def from_json(cls, prediction_from_file: dict) -> "OverallFilePredictions":

tests/test_predictions.py

@@ -57,9 +57,9 @@ def test_to_json(sample_file_prediction):
     result = sample_file_prediction.to_json()
 
     assert isinstance(result, dict)
-    assert result["test_file"]["file_name"] == "test_file"
-    assert len(result["test_file"]["layers"]) == 2
-    assert result["test_file"]["metadata"]["coordinates"]["E"] == 2789456
+    assert result["file_name"] == "test_file"
+    assert len(result["layers"]) == 2
+    assert result["metadata"]["coordinates"]["E"] == 2789456
 
 
 def test_count_against_ground_truth():
@@ -76,13 +76,13 @@ def test_count_against_ground_truth():
 def test_overall_file_predictions():
     """Test OverallFilePredictions class functionality."""
     overall_predictions = OverallFilePredictions()
-    file_prediction = Mock(to_json=lambda: {"test_file": "some_data"})
+    file_prediction = Mock(to_json=lambda: {"some_data": "test"}, file_name="test_file")
 
     overall_predictions.add_file_predictions(file_prediction)
     result = overall_predictions.to_json()
 
     assert len(result) == 1
-    assert result == {"test_file": "some_data"}
+    assert result == {"test_file": {"some_data": "test"}}
 
 
 def test_evaluate_groundwater(sample_file_prediction):