Add the right return type for Language.pipe and an overload for the as_tuples case (#8441)

* Add the right return type for Language.pipe and an overload for the as_tuples version

* Reformat, tidy up

Co-authored-by: Adriane Boyd <adrianeboyd@gmail.com>
This commit is contained in:
Luca Dorigo 2021-07-06 14:18:40 +02:00 committed by GitHub
parent b9f59118bf
commit e8ef4a46d5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1,4 +1,5 @@
from typing import Optional, Any, Dict, Callable, Iterable, Union, List, Pattern from typing import Iterator, Optional, Any, Dict, Callable, Iterable, TypeVar
from typing import Union, List, Pattern, overload
from typing import Tuple from typing import Tuple
from dataclasses import dataclass from dataclasses import dataclass
import random import random
@ -1431,6 +1432,21 @@ class Language:
except StopIteration: except StopIteration:
pass pass
_AnyContext = TypeVar("_AnyContext")
@overload
def pipe(
self,
texts: Iterable[Tuple[str, _AnyContext]],
*,
as_tuples: bool = ...,
batch_size: Optional[int] = ...,
disable: Iterable[str] = ...,
component_cfg: Optional[Dict[str, Dict[str, Any]]] = ...,
n_process: int = ...,
) -> Iterator[Tuple[Doc, _AnyContext]]:
...
def pipe( def pipe(
self, self,
texts: Iterable[str], texts: Iterable[str],
@ -1440,7 +1456,7 @@ class Language:
disable: Iterable[str] = SimpleFrozenList(), disable: Iterable[str] = SimpleFrozenList(),
component_cfg: Optional[Dict[str, Dict[str, Any]]] = None, component_cfg: Optional[Dict[str, Dict[str, Any]]] = None,
n_process: int = 1, n_process: int = 1,
): ) -> Iterator[Doc]:
"""Process texts as a stream, and yield `Doc` objects in order. """Process texts as a stream, and yield `Doc` objects in order.
texts (Iterable[str]): A sequence of texts to process. texts (Iterable[str]): A sequence of texts to process.