Fix return type for mypy

This commit is contained in:
thomashacker 2022-09-05 14:56:33 +02:00
parent 0adc7fde54
commit 064ba9d2bc

View File

@ -98,9 +98,30 @@ class Underscore:
def has(self, name: str) -> bool:
return name in self._extensions
def _get_key(self, name: str) -> Tuple[str, str, Optional[int], Optional[int]]:
def _get_key(
self, name: str
) -> Union[
Tuple[str, str, Optional[int], Optional[int]],
Tuple[
str,
str,
Optional[int],
Optional[int],
Optional[Union[str, int]],
Optional[Union[str, int]],
Optional[Union[str, int]],
],
]:
if hasattr(self, "_label"):
return ("._.", name, self._start, self._end, self._label, self._kb_id, self._span_id)
return (
"._.",
name,
self._start,
self._end,
self._label,
self._kb_id,
self._span_id,
)
else:
return ("._.", name, self._start, self._end)
@ -114,7 +135,9 @@ class Underscore:
old_key = old_underscore._get_key(name)
new_key = new_underscore._get_key(name)
if old_key in old_underscore._doc.user_data and old_key != new_key:
old_underscore._doc.user_data[new_key] = old_underscore._doc.user_data.pop(old_key)
old_underscore._doc.user_data[
new_key
] = old_underscore._doc.user_data.pop(old_key)
@classmethod
def get_state(cls) -> Tuple[Dict[Any, Any], Dict[Any, Any], Dict[Any, Any]]: