Add types to our context variables

shinthor · Jul 26, 2023 · 7f46abc · 7f46abc
1 parent dc597d2
commit 7f46abc
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/chainlit/context.py b/src/chainlit/context.py
@@ -1,5 +1,5 @@
-import contextvars
 from asyncio import AbstractEventLoop
+from contextvars import ContextVar
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -11,8 +11,8 @@ def __init__(self, msg="Chainlit context not found", *args, **kwargs):
         super().__init__(msg, *args, **kwargs)
 
 
-emitter_var = contextvars.ContextVar("emitter")
-loop_var = contextvars.ContextVar("loop")
+emitter_var: ContextVar["ChainlitEmitter"] = ContextVar("emitter")
+loop_var: ContextVar[AbstractEventLoop] = ContextVar("loop")
 
 
 def get_emitter() -> "ChainlitEmitter":