diff --git a/requirements.txt b/requirements.txt index 4394738..33072ee 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,7 +11,7 @@ mock isort black flake8 -mypy<2.0.0 +mypy # Distribution setuptools; python_version >= "3.12" diff --git a/time_execution/timed.py b/time_execution/timed.py index cc78d5d..2db31bf 100644 --- a/time_execution/timed.py +++ b/time_execution/timed.py @@ -62,7 +62,8 @@ def enter(self) -> Any: self._start_time = default_timer() for hook in self._hooks: if isgenerator(hook): - hook.send(None) # start a generator hook + gen = cast(GeneratorHookReturnType, hook) + next(gen) # advance to the first yield return self def get_metric(self) -> Dict[str, Any]: @@ -94,8 +95,9 @@ def apply_hook( ) else: # Generator hook: send the results and obtain custom metadata. + gen = cast(GeneratorHookReturnType, hook) try: - hook.send((self.result, exception, metric)) + gen.send((self.result, exception, metric)) except StopIteration as e: hook_result = e.value else: