Improving type system for pyperf library

pyperf/_cpu_utils.py

@@ -1,3 +1,5 @@
+from typing import Literal
+from collections.abc import Sequence
 import collections
 import os
 import re
@@ -6,14 +8,13 @@
 
 try:
     if not USE_PSUTIL:
-        psutil = None
-    else:
-        import psutil
+        raise ImportError
+    import psutil
 except ImportError:
-    psutil = None
+    psutil = None  # type: ignore[assignment]
 
 
-def get_logical_cpu_count():
+def get_logical_cpu_count() -> int | None:
     if psutil is not None:
         # Number of logical CPUs
         cpu_count = psutil.cpu_count()
@@ -28,15 +29,15 @@ def get_logical_cpu_count():
     return cpu_count
 
 
-def format_cpu_list(cpus):
+def format_cpu_list(cpus: Sequence[int]) -> str:
+    if len(cpus)==0:
+        return "None"
     cpus = sorted(cpus)
     parts = []
-    first = None
-    last = None
-    for cpu in cpus:
-        if first is None:
-            first = cpu
-        elif cpu != last + 1:
+    first = cpus[0]
+    last = cpus[0]
+    for cpu in cpus[1:]:
+        if cpu != last + 1:
             if first != last:
                 parts.append('%s-%s' % (first, last))
             else:
@@ -50,7 +51,7 @@ def format_cpu_list(cpus):
     return ','.join(parts)
 
 
-def format_cpu_infos(infos):
+def format_cpu_infos(infos: dict[int, str]) -> list[str]:
     groups = collections.defaultdict(list)
     for cpu, info in infos.items():
         groups[info].append(cpu)
@@ -59,19 +60,19 @@ def format_cpu_infos(infos):
     items.sort()
     text = []
     for cpus, info in items:
-        cpus = format_cpu_list(cpus)
-        text.append('%s=%s' % (cpus, info))
+        cpu_str = format_cpu_list(cpus)
+        text.append('%s=%s' % (cpu_str, info))
     return text
 
 
-def parse_cpu_list(cpu_list):
+def parse_cpu_list(cpu_list: str) -> list[int] | None:
     cpu_list = cpu_list.strip(' \x00')
     # /sys/devices/system/cpu/nohz_full returns ' (null)\n' when NOHZ full
     # is not used
     if cpu_list == '(null)':
-        return
+        return None
     if not cpu_list:
-        return
+        return None
 
     cpus = []
     for part in cpu_list.split(','):
@@ -88,15 +89,15 @@ def parse_cpu_list(cpu_list):
     return cpus
 
 
-def parse_cpu_mask(line):
+def parse_cpu_mask(line: str) -> int:
     mask = 0
     for part in line.split(','):
         mask <<= 32
         mask |= int(part, 16)
     return mask
 
 
-def format_cpu_mask(mask):
+def format_cpu_mask(mask: int) -> str:
     parts = []
     while 1:
         part = "%08x" % (mask & 0xffffffff)
@@ -107,14 +108,14 @@ def format_cpu_mask(mask):
     return ','.join(reversed(parts))
 
 
-def format_cpus_as_mask(cpus):
+def format_cpus_as_mask(cpus: Sequence[int]) -> str:
     mask = 0
     for cpu in cpus:
         mask |= (1 << cpu)
     return format_cpu_mask(mask)
 
 
-def get_isolated_cpus():
+def get_isolated_cpus() -> list[int] | None:
     """Get the list of isolated CPUs.
 
     Return a sorted list of CPU identifiers, or return None if no CPU is
@@ -137,48 +138,47 @@ def get_isolated_cpus():
     return None
 
 
-def set_cpu_affinity(cpus):
+def set_cpu_affinity(cpus: list[int]) -> Literal[True] | None:
     # Availability: some Unix platforms
     if hasattr(os, 'sched_setaffinity'):
         os.sched_setaffinity(0, cpus)
         return True
 
     try:
         if not USE_PSUTIL:
-            return
-        else:
-            import psutil
+            raise ImportError
+        import psutil
     except ImportError:
-        return
+        return None
 
     # Availability: Linux, Windows, FreeBSD (psutil 2.2.0+)
     # https://psutil.rtfd.io/en/latest/index.html#psutil.Process.cpu_affinity
     proc = psutil.Process()
     if not hasattr(proc, 'cpu_affinity'):
-        return
+        return None
 
     proc.cpu_affinity(cpus)
     return True
 
 
-def set_highest_priority():
+def set_highest_priority() -> Literal[True] | None:
     try:
         if not USE_PSUTIL:
-            return
-        else:
-            import psutil
+            raise ImportError
+        import psutil
     except ImportError:
-        return
+        return None
 
     proc = psutil.Process()
     if not hasattr(proc, 'nice'):
-        return
+        return None
 
     # Want to set realtime on Windows.
     # Fail hard for anything else right now, so it is obvious what to fix
     # when adding other OS support.
     try:
-        proc.nice(psutil.REALTIME_PRIORITY_CLASS)
+        proc.nice(psutil.REALTIME_PRIORITY_CLASS) # type: ignore[attr-defined,unused-ignore]
         return True
     except psutil.AccessDenied:
         pass
+    return None

pyperf/_formatter.py

@@ -1,7 +1,10 @@
+from typing import Literal, cast
+from collections.abc import Sequence
+import datetime
 _TIMEDELTA_UNITS = ('sec', 'ms', 'us', 'ns')
 
 
-def format_timedeltas(values):
+def format_timedeltas(values: Sequence[float]) -> tuple[str, ...]:
     ref_value = abs(values[0])
     for i in range(2, -9, -1):
         if ref_value >= 10.0 ** i:
@@ -18,11 +21,11 @@ def format_timedeltas(values):
     return tuple(fmt % (value * factor,) for value in values)
 
 
-def format_timedelta(value):
+def format_timedelta(value: float) -> str:
     return format_timedeltas((value,))[0]
 
 
-def format_filesize(size):
+def format_filesize(size: float) -> str:
     if size < 10 * 1024:
         if size != 1:
             return '%.0f bytes' % size
@@ -35,11 +38,11 @@ def format_filesize(size):
     return '%.1f KiB' % (size / 1024.0)
 
 
-def format_filesizes(sizes):
+def format_filesizes(sizes: Sequence[float]) -> tuple[str, ...]:
     return tuple(format_filesize(size) for size in sizes)
 
 
-def format_seconds(seconds):
+def format_seconds(seconds: float | None) -> str:
     # Coarse but human readable duration
     if not seconds:
         return '0 sec'
@@ -64,8 +67,9 @@ def format_seconds(seconds):
     return ' '.join(parts)
 
 
-def format_number(number, unit=None, units=None):
+def format_number(number: int, unit: str | None = None, units: str | None = None) -> str:
     plural = (not number or abs(number) > 1)
+    str_number = str(number)
     if number >= 10000:
         pow10 = 0
         x = number
@@ -75,47 +79,51 @@ def format_number(number, unit=None, units=None):
             if r:
                 break
         if not r:
-            number = '10^%s' % pow10
+            str_number = '10^%s' % pow10
 
-    if isinstance(number, int) and number > 8192:
+    if str_number == str(number) and number > 8192:
         if number & (number-1) == 0:
-            number = '2^%s' % (number.bit_length() - 1)
+            str_number = '2^%s' % (number.bit_length() - 1)
 
     if not unit:
-        return str(number)
+        return str_number
 
     if plural:
         if not units:
             units = unit + 's'
-        return '%s %s' % (number, units)
+        return '%s %s' % (str_number, units)
     else:
-        return '%s %s' % (number, unit)
+        return '%s %s' % (str_number, unit)
 
 
-def format_integers(numbers):
+def format_integers(numbers: Sequence[int]) -> tuple[str, ...]:
     return tuple(format_number(number) for number in numbers)
 
 
-DEFAULT_UNIT = 'second'
+
 UNIT_FORMATTERS = {
     'second': format_timedeltas,
     'byte': format_filesizes,
     'integer': format_integers,
 }
-
-
-def format_values(unit, values):
+_UNIT_TYPE = Literal['second', 'byte', 'integer']
+DEFAULT_UNIT: _UNIT_TYPE = 'second'
+def format_values(unit: _UNIT_TYPE | None, values: Sequence[float | int]) -> tuple[str, ...]:
     if not unit:
         unit = DEFAULT_UNIT
-    formatter = UNIT_FORMATTERS[unit]
-    return formatter(values)
+    if unit == 'second':
+        return format_timedeltas(values)
+    elif unit == 'byte':
+        return format_filesizes(values)
+    elif unit == 'integer':
+        return format_integers(cast(Sequence[int],values))
 
 
-def format_value(unit, value):
+def format_value(unit: _UNIT_TYPE, value: float | int) -> str:
     return format_values(unit, (value,))[0]
 
 
-def format_datetime(dt, microsecond=True):
+def format_datetime(dt: datetime.datetime, microsecond: bool = True) -> str:
     if not microsecond:
         dt = dt.replace(microsecond=0)
     return dt.isoformat(' ')