1 files changed, 5 insertions, 6 deletions
diff --git a/poezio/decorators.py b/poezio/decorators.py
index a95e7348..ee138744 100644
--- a/poezio/decorators.py
+++ b/poezio/decorators.py
@@ -13,20 +13,19 @@ from typing import (
     List,
     Optional,
     TypeVar,
-    TYPE_CHECKING,
 )
 
 from poezio import common
 
-if TYPE_CHECKING:
-    from poezio.tabs import RosterInfoTab
 
 T = TypeVar('T', bound=Callable[..., Any])
 
-BeforeFunc = Callable[[List[Any], Dict[str, Any]], Any]
-AfterFunc = Callable[[List[Any], Dict[str, Any]], Any]
 
-def wrap_generic(func: Callable, before: BeforeFunc=None, after: AfterFunc=None):
+BeforeFunc = Optional[Callable[[List[Any], Dict[str, Any]], Any]]
+AfterFunc = Optional[Callable[[Any, List[Any], Dict[str, Any]], Any]]
+
+
+def wrap_generic(func: Callable, before: BeforeFunc = None, after: AfterFunc = None):
     """
     Generic wrapper which can both wrap coroutines and normal functions.
     """