Add typehints for step func returning types

allure-python-commons/src/_allure.py

@@ -1,5 +1,5 @@
 from functools import wraps
-from typing import Any, Callable, TypeVar
+from typing import Any, Callable, TypeVar, Union, overload
 
 from allure_commons._core import plugin_manager
 from allure_commons.types import LabelType, LinkType, ParameterMode
@@ -161,7 +161,15 @@ def manual():
         return Dynamic.label(LabelType.MANUAL, True)
 
 
-def step(title):
+@overload
+def step(title: str) -> "StepContext": ...
+
+
+@overload
+def step(title: _TFunc) -> _TFunc: ...
+
+
+def step(title: Union[str, _TFunc]) -> Union["StepContext", _TFunc]:
     if callable(title):
         return StepContext(title.__name__, {})(title)
     else: