From 25b5de9e2dc6c0542c8b35429e915ad79d30c68b Mon Sep 17 00:00:00 2001
From: "WQY\\qiong" <qiongyu.wang@outlook.com>
Date: Sat, 19 Nov 2022 18:16:52 +0800
Subject: [PATCH] Code refactor

---
 epanet/epanet.py | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/epanet/epanet.py b/epanet/epanet.py
index c6ec611..bcfa732 100644
--- a/epanet/epanet.py
+++ b/epanet/epanet.py
@@ -128,6 +128,20 @@ class Output:
         return usages
 
 
+    def dump(self, cache: bool = False) -> str:
+        data = {}
+        data |= { 'version' : self.version() }
+        data |= { 'net_size' : self.net_size() }
+        data |= { 'units' : self.units() }
+        data |= { 'times' : self.times() }
+        data |= { 'element_name' : self.element_name() }
+        data |= { 'energy_usage' : self.energy_usage() }
+        if cache:
+            with open(self.opt + '.json', 'w') as f:
+                json.dump(data, f)
+        return json.dumps(data)
+
+
     def __del__(self):
         # throw exception in destructor ? :)
         self._check(self._lib.ENR_close(ctypes.byref(self._handle)))
@@ -135,16 +149,7 @@ class Output:
 
 def _read_output(name: str) -> str:
     opt = Output(name)
-    data = {}
-    data |= { 'version' : opt.version() }
-    data |= { 'net_size' : opt.net_size() }
-    data |= { 'units' : opt.units() }
-    data |= { 'times' : opt.times() }
-    data |= { 'element_name' : opt.element_name() }
-    data |= { 'energy_usage' : opt.energy_usage() }
-    with open(opt.opt + '.json', 'w') as f:
-        json.dump(data, f)
-    return json.dumps(data)
+    return opt.dump()
 
 
 def run_project(name: str) -> str: