Avoid counting disk io twice for NVMe SSD block device

The first NVMe block device will be named as nvme0n1,
and its partitions are nvme0n1p1, nvme0n1p2, etc.
Therefore we could not use the last letter to identify if the name is a
partition or not.

psutil/_pslinux.py

@@ -1017,23 +1017,24 @@ def disk_io_counters():
     system as a dict of raw tuples.
     """
     # determine partitions we want to look for
+    partition_suffix = re.compile(r'^p*\d+$')
+
     def get_partitions():
         partitions = []
         with open_text("%s/partitions" % get_procfs_path()) as f:
             lines = f.readlines()[2:]
         for line in reversed(lines):
             _, _, _, name = line.split()
-            if name[-1].isdigit():
-                # we're dealing with a partition (e.g. 'sda1'); 'sda' will
-                # also be around but we want to omit it
-                partitions.append(name)
-            else:
-                if not partitions or not partitions[-1].startswith(name):
-                    # we're dealing with a disk entity for which no
-                    # partitions have been defined (e.g. 'sda' but
-                    # 'sda1' was not around), see:
-                    # https://github.com/giampaolo/psutil/issues/338
+            if partitions and partitions[-1].startswith(name):
+                suffix = partitions[-1][len(name):]
+                if not partition_suffix.match(suffix):
+                    # If a disk entity (e.g. 'sda' or 'nvme0n1') has
+                    # partition(e.g. 'sda1' or 'nvme0n1p1'),
+                    # we will deal with the partition and ignore the
+                    # disk entiy.
                     partitions.append(name)
+            else:
+                partitions.append(name)
         return partitions
 
     retdict = {}

psutil/tests/test_linux.py

@@ -1014,6 +1014,39 @@ def open_mock(name, *args, **kwargs):
             self.assertEqual(ret.write_time, 0)
             self.assertEqual(ret.busy_time, 0)
 
+    def test_disk_io_counters_with_nvme(self):
+        def open_mock(name, *args, **kwargs):
+            if name == '/proc/partitions':
+                return io.StringIO(textwrap.dedent(u"""\
+                    major minor  #blocks  name
+
+                    259        0  250059096 nvme0n1
+                    259        1     266240 nvme0n1p1
+                    259        2      16384 nvme0n1p2
+                    """))
+            elif name == '/proc/diskstats':
+                return io.StringIO(textwrap.dedent(u"""\
+                        259       0 nvme0n1  3 6 9 12 15 18 21 24 27 30 33
+                        259       1 nvme0n1p1 1 2 3 4 5 6 7 8 9 10 11
+                        259       2 nvme0n1p2 2 4 6 8 10 12 14 16 18 20 22
+                    """))
+            else:
+                return orig_open(name, *args, **kwargs)
+
+        orig_open = open
+        patch_point = 'builtins.open' if PY3 else '__builtin__.open'
+        with mock.patch(patch_point, side_effect=open_mock) as m:
+            ret = psutil.disk_io_counters(nowrap=False)
+            assert m.called
+            self.assertEqual(ret.read_count, 3)
+            self.assertEqual(ret.read_merged_count, 6)
+            self.assertEqual(ret.read_bytes, 9 * SECTOR_SIZE)
+            self.assertEqual(ret.read_time, 12)
+            self.assertEqual(ret.write_count, 15)
+            self.assertEqual(ret.write_merged_count, 18)
+            self.assertEqual(ret.write_bytes, 21 * SECTOR_SIZE)
+            self.assertEqual(ret.write_time, 24)
+            self.assertEqual(ret.busy_time, 30)
 
 # =====================================================================
 # --- misc