typehints for Vocab __init__

Author: leezu

src/gluonnlp/vocab/vocab.py

@@ -20,21 +20,19 @@
 # pylint: disable=consider-iterating-dictionary
 
 """Vocabulary."""
-from __future__ import absolute_import
-from __future__ import print_function
-
 __all__ = ['Vocab']
 
 import collections
 import json
+import typing
 import uuid
 import warnings
 
 from mxnet import nd
 
-from ..data.utils import DefaultLookupDict, count_tokens
 from .. import _constants as C
 from .. import embedding as emb
+from ..data.utils import DefaultLookupDict, Counter, count_tokens
 
 UNK_IDX = 0
 
@@ -44,38 +42,38 @@ class Vocab(object):
 
     Parameters
     ----------
-    counter : Counter or None, default None
+    counter
         Counts text token frequencies in the text data. Its keys will be indexed according to
         frequency thresholds such as `max_size` and `min_freq`. Keys of `counter`,
         `unknown_token`, and values of `reserved_tokens` must be of the same hashable type.
         Examples: str, int, and tuple.
-    max_size : None or int, default None
+    max_size
         The maximum possible number of the most frequent tokens in the keys of `counter` that can be
         indexed. Note that this argument does not count any token from `reserved_tokens`. Suppose
         that there are different keys of `counter` whose frequency are the same, if indexing all of
         them will exceed this argument value, such keys will be indexed one by one according to
         their __cmp__() order until the frequency threshold is met. If this argument is None or
         larger than its largest possible value restricted by `counter` and `reserved_tokens`, this
         argument has no effect.
-    min_freq : int, default 1
+    min_freq
         The minimum frequency required for a token in the keys of `counter` to be indexed.
-    unknown_token : hashable object or None, default '<unk>'
+    unknown_token
         The representation for any unknown token. If `unknown_token` is not
         `None`, looking up any token that is not part of the vocabulary and
         thus considered unknown will return the index of `unknown_token`. If
         None, looking up an unknown token will result in `KeyError`.
-    padding_token : hashable object or None, default '<pad>'
+    padding_token
         The representation for the special token of padding token.
-    bos_token : hashable object or None, default '<bos>'
+    bos_token
         The representation for the special token of beginning-of-sequence token.
-    eos_token : hashable object or None, default '<eos>'
+    eos_token
         The representation for the special token of end-of-sequence token.
-    reserved_tokens : list of hashable objects or None, default None
+    reserved_tokens
         A list specifying additional tokens to be added to the vocabulary.
         `reserved_tokens` must not contain the value of `unknown_token` or
         duplicate tokens. It must neither contain special tokens specified via
         keyword arguments.
-    token_to_idx : dict mapping tokens (hashable objects) to int or None, default None
+    token_to_idx
         If not `None`, specifies the indices of tokens to be used by the
         vocabulary. Each token in `token_to_index` must be part of the Vocab
         and each index can only be associated with a single token.
@@ -175,9 +173,14 @@ class Vocab(object):
 
     """
 
-    def __init__(self, counter=None, max_size=None, min_freq=1, unknown_token=C.UNK_TOKEN,
-                 padding_token=C.PAD_TOKEN, bos_token=C.BOS_TOKEN, eos_token=C.EOS_TOKEN,
-                 reserved_tokens=None, token_to_idx=None, **kwargs):
+    def __init__(self, counter: typing.Optional[Counter] = None,
+                 max_size: typing.Optional[int] = None, min_freq: int = 1,
+                 unknown_token: typing.Optional[typing.Hashable] = C.UNK_TOKEN,
+                 padding_token: typing.Optional[typing.Hashable] = C.PAD_TOKEN,
+                 bos_token: typing.Optional[typing.Hashable] = C.BOS_TOKEN,
+                 eos_token: typing.Optional[typing.Hashable] = C.EOS_TOKEN,
+                 reserved_tokens: typing.Optional[typing.List[typing.Hashable]] = None,
+                 token_to_idx: typing.Optional[typing.Dict[typing.Hashable, int]] = None, **kwargs):
 
         # Sanity checks.
         assert min_freq > 0, '`min_freq` must be set to a positive value.'