realpython · gahjelle · Nov 8, 2022 · Nov 8, 2022 · Nov 8, 2022
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ### Changed
 
+- Add type hints to tests ([#50])
 - Drop explicit support for Python 3.6 ([#48])
 
 ## [1.4.0] - 2022-11-08
@@ -77,6 +78,7 @@ Initial version of `codetiming`. Version 1.0.0 corresponds to the code in the tu
 [1.1.0]: https://github.com/realpython/codetiming/compare/v1.0.0...v1.1.0
 [1.0.0]: https://github.com/realpython/codetiming/releases/tag/v1.0.0
 
+[#50]: https://github.com/realpython/codetiming/pull/50
 [#48]: https://github.com/realpython/codetiming/pull/48
 [#47]: https://github.com/realpython/codetiming/pull/47
 [#46]: https://github.com/realpython/codetiming/pull/46

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -50,7 +50,7 @@ Run tests using [`tox`](https://tox.readthedocs.io/). You can also run individua
 - Static type hinting using [`mypy`](http://mypy-lang.org/). Test your type hints as follows:
 
     ```console
-    $ mypy --strict codetiming/
+    $ mypy --strict codetiming/ tests/
     ```
 
     See Real Python's [Python Type Checking guide](https://realpython.com/python-type-checking/) for more information.

diff --git a/codetiming/_timer.py b/codetiming/_timer.py
@@ -9,12 +9,32 @@
 import time
 from contextlib import ContextDecorator
 from dataclasses import dataclass, field
-from typing import Any, Callable, ClassVar, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Union
 
 # Codetiming imports
 from codetiming._timers import Timers
 
+# Special types, Protocol only works for Python >= 3.8
+if TYPE_CHECKING:  # pragma: nocover
+    # Standard library imports
+    from typing import Protocol, TypeVar
 
+    T = TypeVar("T")
+
+    class FloatArg(Protocol):
+        """Protocol type that allows classes that take one float argument"""
+
+        def __call__(self: T, __seconds: float) -> T:
+            """Callable signature"""
+            ...  # pragma: nocover
+
+else:
+
+    class FloatArg:
+        """Dummy runtime class"""
+
+
+# Timer code
 class TimerError(Exception):
     """A custom exception used to report errors in use of Timer class."""
 
@@ -24,12 +44,12 @@ class Timer(ContextDecorator):
     """Time your code using a class, context manager, or decorator."""
 
     timers: ClassVar[Timers] = Timers()
-    _start_time: Optional[float] = field(default=None, init=False, repr=False)
     name: Optional[str] = None
-    text: Union[str, Callable[[float], str]] = "Elapsed time: {:0.4f} seconds"
-    initial_text: Union[bool, str] = False
+    text: Union[str, FloatArg, Callable[[float], str]] = "Elapsed time: {:0.4f} seconds"
+    initial_text: Union[bool, str, FloatArg] = False
     logger: Optional[Callable[[str], None]] = print
     last: float = field(default=math.nan, init=False, repr=False)
+    _start_time: Optional[float] = field(default=None, init=False, repr=False)
 
     def start(self) -> None:
         """Start a new timer."""
@@ -69,7 +89,7 @@ def stop(self) -> float:
                     "minutes": self.last / 60,
                 }
                 text = self.text.format(self.last, **attributes)
-            self.logger(text)
+            self.logger(str(text))
         if self.name:
             self.timers.add(self.name, self.last)