Detail type hints

yzhangcs · yzhangcs · commit 7f36190b8a64 · 2022-06-14T16:03:38.000+08:00
diff --git a/supar/utils/vocab.py b/supar/utils/vocab.py
@@ -4,7 +4,7 @@
 
 from collections import Counter, defaultdict
 from collections.abc import Iterable
-from typing import Tuple
+from typing import Tuple, Union
 
 
 class Vocab(object):
@@ -40,7 +40,7 @@ def __init__(self, counter: Counter, min_freq: int = 1, specials: Tuple = tuple(
     def __len__(self):
         return len(self.itos)
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: Union[int, str, Iterable]) -> Union[str, int, Iterable]:
         if isinstance(key, str):
             return self.stoi[key]
         elif not isinstance(key, Iterable):
