From a998f44b930d5cb04d42d4ae600266aa54913f9f Mon Sep 17 00:00:00 2001
From: Tom Hodson <thomas.hodson@ecmwf.int>
Date: Thu, 5 Dec 2024 15:45:43 +0000
Subject: [PATCH] Update tree

---
 tree_compresser/rust_src/tree.rs         | 14 +++-----------
 tree_compresser/tests/open_climate_dt.py |  6 ++++--
 2 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/tree_compresser/rust_src/tree.rs b/tree_compresser/rust_src/tree.rs
index c9ccd0d..59e6695 100644
--- a/tree_compresser/rust_src/tree.rs
+++ b/tree_compresser/rust_src/tree.rs
@@ -72,17 +72,9 @@ impl TreeNode {
     pub fn to_py_dict(&self, py: Python) -> PyResult<PyObject> {
         let py_dict = PyDict::new(py);
 
-        let formatted_key = format!("{}={}", self.key.key, self.key.value);
-
-        if self.children.is_empty() {
-            py_dict.set_item(formatted_key, PyDict::new(py))?;
-        } else {
-            let children_dict = PyDict::new(py);
-            for child in &self.children {
-                let child_key = format!("{}={}", child.key.key, child.key.value);
-                children_dict.set_item(child_key, child.to_py_dict(py)?)?;
-            }
-            py_dict.set_item(formatted_key, children_dict)?;
+        for child in &self.children {
+            let child_key = format!("{}={}", child.key.key, child.key.value);
+            py_dict.set_item(child_key, child.to_py_dict(py)?)?;
         }
 
         Ok(py_dict.to_object(py))
diff --git a/tree_compresser/tests/open_climate_dt.py b/tree_compresser/tests/open_climate_dt.py
index 501e7dd..e4c9a42 100644
--- a/tree_compresser/tests/open_climate_dt.py
+++ b/tree_compresser/tests/open_climate_dt.py
@@ -1,5 +1,6 @@
 from tree_traverser import backend, CompressedTree
 from pathlib import Path
+import json
 
 data_path = Path("data/compressed_tree_climate_dt.json")
 # Print size of file
@@ -8,5 +9,6 @@ print(f"climate dt compressed tree: {data_path.stat().st_size // 1e6:.1f} MB")
 print("Opening json file")
 compressed_tree = CompressedTree.load(data_path)
 
-print("Printing compressed tree")
-print(compressed_tree.reconstruct_compressed_ecmwf_style())
+print("Outputting compressed tree ecmwf style")
+with open("data/compressed_tree_climate_dt_ecmwf_style.json", "w") as f:
+    json.dump(compressed_tree.reconstruct_compressed_ecmwf_style(), f)