From 0a92d5644ef3fcb201f284e623ada019c4ab0b8b Mon Sep 17 00:00:00 2001 From: Luca Dorigo Date: Tue, 3 May 2022 17:57:07 +0200 Subject: [PATCH] Fix StringStore.__getitem__ return type depending on parameter types (#10741) * Fix StringStore.__getitem__ return type depending on parameter types Small fix using `@overload` so that `StringStore.__getitem__` returns an `int` when given a `str` or `bytes` and a `str` when given an `int`. * Update spacy/strings.pyi Co-authored-by: Adriane Boyd --- spacy/strings.pyi | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/spacy/strings.pyi b/spacy/strings.pyi index 5b4147e12..b29389b9a 100644 --- a/spacy/strings.pyi +++ b/spacy/strings.pyi @@ -1,4 +1,4 @@ -from typing import Optional, Iterable, Iterator, Union, Any +from typing import Optional, Iterable, Iterator, Union, Any, overload from pathlib import Path def get_string_id(key: Union[str, int]) -> int: ... @@ -7,7 +7,10 @@ class StringStore: def __init__( self, strings: Optional[Iterable[str]] = ..., freeze: bool = ... ) -> None: ... - def __getitem__(self, string_or_id: Union[bytes, str, int]) -> Union[str, int]: ... + @overload + def __getitem__(self, string_or_id: Union[bytes, str]) -> int: ... + @overload + def __getitem__(self, string_or_id: int) -> str: ... def as_int(self, key: Union[bytes, str, int]) -> int: ... def as_string(self, key: Union[bytes, str, int]) -> str: ... def add(self, string: str) -> int: ...