Unverified Commit 003c4270 authored by Tanmay Verma's avatar Tanmay Verma Committed by GitHub
Browse files

fix: Fix mypy errors on trtllm examples (#1277)

parent 3e3c3b10
...@@ -21,7 +21,7 @@ import traceback ...@@ -21,7 +21,7 @@ import traceback
import weakref import weakref
from enum import Enum from enum import Enum
from queue import Queue from queue import Queue
from typing import Callable, Optional, TypedDict, Union from typing import Any, Callable, Coroutine, Optional, TypedDict, Union
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
...@@ -54,7 +54,9 @@ class ConversationMessage(TypedDict): ...@@ -54,7 +54,9 @@ class ConversationMessage(TypedDict):
class ManagedThread(threading.Thread): class ManagedThread(threading.Thread):
def __init__( def __init__(
self, self,
task: Optional[Union[Callable[..., bool], weakref.WeakMethod]], task: Optional[
Union[Callable[..., Coroutine[Any, Any, bool]], weakref.WeakMethod]
],
error_queue: Optional[Queue] = None, error_queue: Optional[Queue] = None,
name: Optional[str] = None, name: Optional[str] = None,
loop: Optional[asyncio.AbstractEventLoop] = None, loop: Optional[asyncio.AbstractEventLoop] = None,
...@@ -74,7 +76,9 @@ class ManagedThread(threading.Thread): ...@@ -74,7 +76,9 @@ class ManagedThread(threading.Thread):
def run(self): def run(self):
while not self.stop_event.is_set(): while not self.stop_event.is_set():
task: Optional[Union[Callable[..., bool], weakref.WeakMethod]] = self.task task: Optional[
Union[Callable[..., Coroutine[Any, Any, bool]], weakref.WeakMethod]
] = self.task
if isinstance(task, weakref.WeakMethod): if isinstance(task, weakref.WeakMethod):
task = task() task = task()
if task is None: if task is None:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment