diff --git a/stubs/assertpy/assertpy/assertpy.pyi b/stubs/assertpy/assertpy/assertpy.pyi
index 5366825bcbd2..922a86dcf426 100644
--- a/stubs/assertpy/assertpy/assertpy.pyi
+++ b/stubs/assertpy/assertpy/assertpy.pyi
@@ -69,5 +69,5 @@ def assert_that(val: _V, description: str = "") -> AssertionBuilder[_V]: ...
 def assert_warn(val: _V, description: str = "", logger: logging.Logger | None = None) -> AssertionBuilder: ...
 def fail(msg: str = "") -> None: ...
 def soft_fail(msg: str = "") -> None: ...
-def add_extension(func: Callable[[AssertionBuilder], AssertionBuilder]) -> None: ...
-def remove_extension(func: Callable[[AssertionBuilder], AssertionBuilder]) -> None: ...
+def add_extension(func: Callable[..., AssertionBuilder[Any]]) -> None: ...
+def remove_extension(func: Callable[..., AssertionBuilder[Any]]) -> None: ...