From 742a717b4922e858d266efb2a2935df84bae40d0 Mon Sep 17 00:00:00 2001
From: zmy <izmy@qq.com>
Date: Tue, 22 Dec 2020 20:17:24 +0800
Subject: [PATCH] :bug: fix bug on empty configuration file

---
 gpuutil/__main__.py     |  5 ++---
 gpuutil/gpuutil.py      | 27 +++++++++++++++++----------
 gpuutil/set_redirect.py |  5 ++---
 3 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/gpuutil/__main__.py b/gpuutil/__main__.py
index 7c2ae8a..e061b69 100644
--- a/gpuutil/__main__.py
+++ b/gpuutil/__main__.py
@@ -1,4 +1,4 @@
-from gpuutil import GPUStat
+from gpuutil import GPUStat, loaddict
 import sys
 import json
 import argparse
@@ -18,8 +18,7 @@ def load_config():
     configpath = os.path.join(home_dir, '.gpuutil.conf')
     if not os.path.isfile(configpath):
         return {}
-    with open(configpath, 'r', encoding='utf-8') as f:
-        return json.load(f)
+    return loaddict(configpath)
 def save_config(config):
     home_dir = os.path.expanduser('~')
     configdir = os.path.join(home_dir, '.gpuutil.conf')
diff --git a/gpuutil/gpuutil.py b/gpuutil/gpuutil.py
index 83cdefc..e48cc85 100644
--- a/gpuutil/gpuutil.py
+++ b/gpuutil/gpuutil.py
@@ -14,6 +14,14 @@ def loadfile(path):
     with open(path, 'r', encoding='utf-8') as f:
         return f.read()
 
+def loaddict(path):
+    content = loadfile(path)
+    content = content.strip()
+    if len(content) != 0:
+        return json.loads(content)
+    else:
+        return {}
+
 def exe_cmd(command):
     pipe = os.popen(command)
     return pipe.read()
@@ -311,16 +319,15 @@ class GPUStat():
         self.load_configure()
     def load_configure(self):
         configuration_path = os.path.expanduser('~/.gpuutil.conf')
-        if os.path.exists(configuration_path):
-            with open(configuration_path, 'r', encoding='utf-8') as f:
-                configuration = json.load(f)
-                if 'redirect' in configuration:
-                    if 'nvsmi_src' in configuration['redirect']:
-                        self.nvsmi_source = configuration['redirect']['nvsmi_src']
-                    if 'ps_src' in configuration['redirect']:
-                        self.ps_source = configuration['redirect']['ps_src']
-                    if 'ps_name_trans' in configuration['redirect']:
-                        self.ps_name_trans = configuration['redirect']['ps_name_trans']
+        if os.path.isfile(configuration_path):
+            configuration = loaddict(configuration_path)
+            if 'redirect' in configuration:
+                if 'nvsmi_src' in configuration['redirect']:
+                    self.nvsmi_source = configuration['redirect']['nvsmi_src']
+                if 'ps_src' in configuration['redirect']:
+                    self.ps_source = configuration['redirect']['ps_src']
+                if 'ps_name_trans' in configuration['redirect']:
+                    self.ps_name_trans = configuration['redirect']['ps_name_trans']
 
             
     def get_process_info(self):
diff --git a/gpuutil/set_redirect.py b/gpuutil/set_redirect.py
index 84af30b..ad5b16c 100644
--- a/gpuutil/set_redirect.py
+++ b/gpuutil/set_redirect.py
@@ -1,6 +1,6 @@
 import argparse
 import os
-import json
+from gpuutil import loaddict
 
 availabel_name_trans = ['command', 'user', 'pid']
 
@@ -32,8 +32,7 @@ if name_trans is not None:
 config_file = os.path.expanduser('~/.gpuutil.conf')
 configuration = {}
 if os.path.isfile(config_file):
-    with open(config_file, 'r', encoding='utf-8') as f:
-        configuration = json.load(f)
+    configuration = loaddict(config_file)
 configuration['redirect'] = {
     "nvsmi_src": args.nvsmi,
     "ps_src": args.ps,