From c90b22f6ce93ac1ca9eaa94ba052197ccde4f9d7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=A4=D0=B5=D0=B4=D0=BE=D1=81=D0=BE=D0=B2=20=D0=9D=D0=B8?=
 =?UTF-8?q?=D0=BA=D0=B8=D1=82=D0=B0?= <n.fedosov@2gis.ru>
Date: Fri, 20 Sep 2024 09:39:00 +0500
Subject: [PATCH] Add typehints for step func returning types

---
 allure-python-commons/src/_allure.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/allure-python-commons/src/_allure.py b/allure-python-commons/src/_allure.py
index 05e01dbd..84bc755f 100644
--- a/allure-python-commons/src/_allure.py
+++ b/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: