#802: handle the case where dict keys changes between calls

2 files changed, 16 insertions, 2 deletions

diff --git a/psutil/_common.py b/psutil/_common.py
index 9e604b6a..202c371a 100644
--- a/psutil/_common.py
+++ b/psutil/_common.py
@@ -484,8 +484,14 @@ def wrap_numbers(input_dict, name):
         new_dict = {}
         old_dict = _wrapn_cache[name]
         for key in input_dict.keys():
-            old_nt = old_dict[key]
             input_nt = input_dict[key]
+            try:
+                old_nt = old_dict[key]
+            except KeyError:
+                # The input dict has a new key (e.g. a new disk or NIC)
+                # which didn't exist in the previous call.
+                new_dict[key] = input_nt
+                continue
 
             bits = []
             for i in range(len(input_nt)):
diff --git a/psutil/tests/test_misc.py b/psutil/tests/test_misc.py
index 795b58e2..799dd6b5 100755
--- a/psutil/tests/test_misc.py
+++ b/psutil/tests/test_misc.py
@@ -410,7 +410,7 @@ class TestWrapNumbers(unittest.TestCase):
         # let's say 100 is the threshold
         input = {'foo': nt(100, 100, 100)}
         self.assertEqual(wrap_numbers(input, 'funname'), input)
-        # first wrap restart from 10
+        # first wrap restarts from 10
         input = {'foo': nt(100, 100, 10)}
         self.assertEqual(wrap_numbers(input, 'funname'),
                          {'foo': nt(100, 100, 110)})
@@ -427,6 +427,14 @@ class TestWrapNumbers(unittest.TestCase):
         self.assertEqual(wrap_numbers(input, 'funname'),
                          {'foo': nt(100, 100, 210)})
 
+    def test_dict_keys_mismatch(self):
+        # Emulate a case where the second call to disk_io_counters()
+        # (or whatever) provides a new disk.
+        input = {'disk1': nt(5, 5, 5)}
+        self.assertEqual(wrap_numbers(input, 'funname'), input)
+        input = {'disk1': nt(5, 5, 5), 'disk2': nt(7, 7, 7)}
+        self.assertEqual(wrap_numbers(input, 'funname'), input)
+
 
 # ===================================================================
 # --- Example script tests