Add type hints for ExpectQ methods

This commit is contained in:
Jiang Yio 2025-01-14 22:42:44 -05:00
parent d261d5200e
commit c0d9cbfc42

View File

@ -9,7 +9,7 @@ import contextlib
import logging
from collections import namedtuple
from typing import Optional, Union, Sequence, NamedTuple, Callable
from typing import Any, Optional, Union, Sequence, NamedTuple, Callable, AsyncGenerator
logger = logging.getLogger(__name__)
@ -28,7 +28,7 @@ class ExpectQ(object):
"""Clear or restore buffer"""
self.buffer = buffer
clear = reset
async def prompts(self, endl: str='\r\n', timeout_settle: Optional[float]=None, throw: bool=False):
async def prompts(self, endl: str='\r\n', timeout_settle: Optional[float]=None, throw: bool=False) -> AsyncGenerator[tuple[Optional[str], Optional[int]], None]:
len_endl = len(endl)
while True:
if (pos := self.buffer.rfind(endl)) >= 0:
@ -43,7 +43,7 @@ class ExpectQ(object):
if throw:
raise
yield None, None
async def promptmatches(self, *mappings: Union[str, re.Pattern, tuple, list], endl: str='\r\n', timeout_settle: Optional[float]=None, throw: bool=False):
async def promptmatches(self, *mappings: Union[str, re.Pattern, tuple, list], endl: str='\r\n', timeout_settle: Optional[float]=None, throw: bool=False) -> AsyncGenerator[tuple[Optional[ExpectMatch], Any], Optional[bool]]:
for i, mapping in enumerate(mappings):
try:
match mapping: