From d2b43e31fded161c4122b1bb61ed1486b2ee4818 Mon Sep 17 00:00:00 2001
From: Geir <geir.drange@gmail.com>
Date: Wed, 22 May 2024 17:03:06 +0200
Subject: [PATCH] fix: print tree in .tree() method to make it show in terminal
 mode

---
 docs/nb/Bluetooth sensor.ipynb | 3 +--
 src/sysloss/system.py          | 8 ++------
 tests/unit/test_system.py      | 6 ++----
 3 files changed, 5 insertions(+), 12 deletions(-)

diff --git a/docs/nb/Bluetooth sensor.ipynb b/docs/nb/Bluetooth sensor.ipynb
index f505fbe..6f47453 100644
--- a/docs/nb/Bluetooth sensor.ipynb	
+++ b/docs/nb/Bluetooth sensor.ipynb	
@@ -131,9 +131,8 @@
        "        └── MCU\n"
       ]
      },
-     "execution_count": 4,
      "metadata": {},
-     "output_type": "execute_result"
+     "output_type": "display_data"
     }
    ],
    "source": [
diff --git a/src/sysloss/system.py b/src/sysloss/system.py
index e460ba3..adef41f 100644
--- a/src/sysloss/system.py
+++ b/src/sysloss/system.py
@@ -26,6 +26,7 @@
 import rustworkx as rx
 import numpy as np
 from rich.tree import Tree
+from rich import print
 import json
 import pandas as pd
 import matplotlib
@@ -436,11 +437,6 @@ def tree(self, name=""):
         name : str, optional
             Name of component to start with. If not given, print the entire system., by default ""
 
-        Returns
-        -------
-        rich.tree.Tree
-            System Tree class.
-
         Raises
         ------
         ValueError
@@ -464,7 +460,7 @@ def tree(self, name=""):
                     c += [j._params["name"]]
                 ndict[i[0]._params["name"]] = c
             t.add(self._make_rtree(ndict, n))
-        return t
+        print(t)
 
     def _set_phase_lkup(self):
         """Make lookup from node # to load phases"""
diff --git a/tests/unit/test_system.py b/tests/unit/test_system.py
index 6251bfa..e1bd11e 100644
--- a/tests/unit/test_system.py
+++ b/tests/unit/test_system.py
@@ -61,12 +61,10 @@ def test_case1():
     case1.save("tests/unit/case1.json")
     dfp = case1.params(limits=True)
     assert len(dfp) == rows - 1, "Case1 parameters row count"
-    t = case1.tree()
-    assert type(t) == rich.tree.Tree, "Case1 tree output"
+    assert case1.tree() == None, "Case1 tree output"
     with pytest.raises(ValueError):
         case1.tree("Dummy")
-    t = case1.tree("5V boost")
-    assert type(t) == rich.tree.Tree, "Case1 subtree output"
+    assert case1.tree("5V boost") == None, "Case1 subtree output"
     edata = {"vi": [3.6], "io": [0.1, 0.4, 0.6, 0.9], "eff": [[0.3, 0.4, 0.67, 0.89]]}
     case1.change_comp(
         "1.8V buck", comp=Converter("1.8V buck", vo=1.8, eff=edata, iq=12e-6)