Fix type hints in Timers class

realpython · gahjelle · Feb 4, 2020 · Jan 27, 2020 · Feb 4, 2020 · Feb 4, 2020
commit e178f8073513168955005975e7a707257fae2d37
diff --git a/codetiming/_timer.py b/codetiming/_timer.py
@@ -9,7 +9,7 @@
 import time
 from contextlib import ContextDecorator
 from dataclasses import dataclass, field
-from typing import Any, Callable, ClassVar, Dict, Optional
+from typing import Any, Callable, ClassVar, Optional
 
 from codetiming._timers import Timers
 
@@ -22,18 +22,13 @@ class TimerError(Exception):
 class Timer(ContextDecorator):
     """Time your code using a class, context manager, or decorator"""
 
-    timers: ClassVar[Dict[str, float]] = Timers()
+    timers: ClassVar[Timers] = Timers()
     _start_time: Optional[float] = field(default=None, init=False, repr=False)
     name: Optional[str] = None
     text: str = "Elapsed time: {:0.4f} seconds"
     logger: Optional[Callable[[str], None]] = print
     last: float = field(default=math.nan, init=False, repr=False)
 
-    def __post_init__(self) -> None:
-        """Initialization: add timer to dict of timers"""
-        # if self.name:
-        #     self.timers.setdefault(self.name, 0)
-
     def start(self) -> None:
         """Start a new timer"""
         if self._start_time is not None:

diff --git a/codetiming/_timers.py b/codetiming/_timers.py
@@ -4,10 +4,16 @@
 import collections
 import math
 import statistics
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, TYPE_CHECKING
 
+# Annotate generic UserDict
+if TYPE_CHECKING:
+    UserDict = collections.UserDict[str, float]  # pragma: no cover
+else:
+    UserDict = collections.UserDict
 
-class Timers(collections.UserDict):
+
+class Timers(UserDict):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Add a private dictionary keeping track of all timings"""
         super().__init__(*args, **kwargs)
@@ -31,47 +37,38 @@ def __setitem__(self, name: str, value: float) -> None:
             "Use '.add()' to update values."
         )
 
-    def apply(
-        self, func: Callable[[List[float]], float], name: Optional[str] = None
-    ) -> float:
-        """Apply function to timings"""
-        if name is None:
-            return {k: func(v) for k, v in self._timings.items()}
+    def apply(self, func: Callable[[List[float]], float], name: str) -> float:
+        """Apply a function to the results of one named timer"""
         if name in self._timings:
             return func(self._timings[name])
         raise KeyError(name)
 
-    def count(self, name: Optional[str] = None) -> float:
+    def count(self, name: str) -> float:
         """Number of timings"""
         return self.apply(len, name=name)
 
-    def total(self, name: Optional[str] = None) -> float:
+    def total(self, name: str) -> float:
         """Total time for timers"""
         return self.apply(sum, name=name)
 
-    def min(self, name=None):
+    def min(self, name: str) -> float:
         """Minimal value of timings"""
         return self.apply(lambda values: min(values or [0]), name=name)
 
-    def max(self, name=None):
+    def max(self, name: str) -> float:
         """Maximal value of timings"""
         return self.apply(lambda values: max(values or [0]), name=name)
 
-    def mean(self, name=None):
+    def mean(self, name: str) -> float:
         """Mean value of timings"""
         return self.apply(lambda values: statistics.mean(values or [0]), name=name)
 
-    def median(self, name=None):
+    def median(self, name: str) -> float:
         """Median value of timings"""
         return self.apply(lambda values: statistics.median(values or [0]), name=name)
 
-    def stdev(self, name=None):
+    def stdev(self, name: str) -> float:
         """Standard deviation of timings"""
-        if name is None:
-            return {
-                k: statistics.stdev(v) if len(v) >= 2 else math.nan
-                for k, v in self._timings.items()
-            }
         if name in self._timings:
             value = self._timings[name]
             return statistics.stdev(value) if len(value) >= 2 else math.nan

diff --git a/tests/test_codetiming.py b/tests/test_codetiming.py
@@ -225,18 +225,6 @@ def test_timers_stats():
     assert stats.stdev(name) >= 0
 
 
-def test_stats_all_timers():
-    """Test that we can get stats dictionaries from timers"""
-    Timer.timers.clear()
-    for num in range(5, 10):
-        with Timer(name=f"timer_{num}"):
-            waste_time(num=100 * num)
-
-    stats = Timer.timers
-    assert isinstance(stats.total(), dict)
-    assert len(stats.stdev()) == 5
-
-
 def test_stats_missing_timers():
     """Test that getting statistics from non-existent timers raises exception"""
     with pytest.raises(KeyError):