Fix missing type hints.

getitem on storage types, various iterators and simplevocab missed
type hints or had incorrect type information.

Author: sebpuetz

mypy.ini

@@ -3,6 +3,10 @@
 [mypy]
 warn_return_any = True
 warn_unused_configs = True
+warn_redundant_casts = True
+warn_unreachable = True
+check_untyped_defs = True
+disallow_untyped_calls = True
 
 # Per-module options:
 

src/finalfusion/_util.py

@@ -2,10 +2,9 @@
 import numpy as np
 
 from finalfusion.norms import Norms
-from finalfusion.storage import NdArray
 
 
-def _normalize_ndarray_storage(storage: NdArray) -> Norms:
+def _normalize_matrix(storage: np.ndarray) -> Norms:
     norms = np.linalg.norm(storage, axis=1)
     storage /= norms[:, None]
     return Norms(norms)

src/finalfusion/compat/fasttext.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 from finalfusion import Embeddings
-from finalfusion._util import _normalize_ndarray_storage
+from finalfusion._util import _normalize_matrix
 from finalfusion.io import _read_required_binary, _write_binary, _serialize_array_as_le
 from finalfusion.metadata import Metadata
 from finalfusion.storage import NdArray
@@ -46,7 +46,7 @@ def load_fasttext(file: Union[str, bytes, int, PathLike]) -> Embeddings:
         vocab = _read_ft_vocab(inf, metadata['buckets'], metadata['min_n'],
                                metadata['max_n'])
         storage = _read_ft_storage(inf, vocab)
-        norms = _normalize_ndarray_storage(storage[:len(vocab)])
+        norms = _normalize_matrix(storage[:len(vocab)])
     return Embeddings(storage=storage,
                       vocab=vocab,
                       norms=norms,
@@ -309,10 +309,12 @@ def _write_ft_storage_subwords(outf: BinaryIO, embeds: Embeddings):
     norms = embeds.norms
     for i, word in enumerate(vocab):
         indices = vocab.subword_indices(word)
-        embed = storage[i] * (len(indices) + 1)
+        embed = storage[i]  # type: np.ndarray
+        embed = embed * (len(indices) + 1)
         if norms is not None:
             embed *= norms[i]
-        embed -= storage[indices].sum(0, keepdims=False)
+        sw_embeds = storage[indices]  # type: np.ndarray
+        embed -= sw_embeds.sum(0, keepdims=False)
         _serialize_array_as_le(outf, embed)
 
     _serialize_array_as_le(outf, storage[len(vocab):])
@@ -327,7 +329,7 @@ def _write_ft_storage_simple(outf: BinaryIO, embeds: Embeddings):
     storage = embeds.storage
     norms = embeds.norms
     for i in range(storage.shape[0]):
-        embed = storage[i]
+        embed = storage[i]  # type: np.ndarray
         if norms is not None:
             embed = norms[i] * embed
         _serialize_array_as_le(outf, embed)

src/finalfusion/compat/text.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 from finalfusion import Embeddings
-from finalfusion._util import _normalize_ndarray_storage
+from finalfusion._util import _normalize_matrix
 from finalfusion.storage import NdArray
 from finalfusion.vocab import SimpleVocab
 
@@ -137,7 +137,7 @@ def _load_text(file: TextIO, rows: int, cols: int) -> Embeddings:
         row[:] = parts[1:]
     storage = NdArray(matrix)
     return Embeddings(storage=storage,
-                      norms=_normalize_ndarray_storage(storage),
+                      norms=_normalize_matrix(storage),
                       vocab=SimpleVocab(words))
 
 
@@ -151,7 +151,7 @@ def _write_text(file: Union[str, bytes, int, PathLike],
         if dims:
             print(*matrix.shape, file=outf)
         for idx, word in enumerate(vocab):
-            row = matrix[idx]
+            row = matrix[idx]  # type: np.ndarray
             if embeddings.norms is not None:
                 row = row * embeddings.norms[idx]
             print(word, ' '.join(map(str, row)), sep=sep, file=outf)

src/finalfusion/compat/word2vec.py

@@ -11,7 +11,7 @@
 from finalfusion import Embeddings
 from finalfusion.io import _serialize_array_as_le
 from finalfusion.storage import NdArray
-from finalfusion._util import _normalize_ndarray_storage
+from finalfusion._util import _normalize_matrix
 from finalfusion.vocab import SimpleVocab
 
 
@@ -48,7 +48,7 @@ def load_word2vec(file: Union[str, bytes, int, PathLike]) -> Embeddings:
             row[:] = array
     storage = NdArray(matrix)
     return Embeddings(storage=storage,
-                      norms=_normalize_ndarray_storage(storage),
+                      norms=_normalize_matrix(storage),
                       vocab=SimpleVocab(words))
 
 
@@ -84,7 +84,7 @@ def write_word2vec(file: Union[str, bytes, int, PathLike],
     with open(file, 'wb') as outf:
         outf.write(f'{matrix.shape[0]} {matrix.shape[1]}\n'.encode('ascii'))
         for idx, word in enumerate(vocab):
-            row = matrix[idx]
+            row = matrix[idx]  # type: np.ndarray
             if embeddings.norms is not None:
                 row = row * embeddings.norms[idx]
             b_word = word.encode('utf-8')

src/finalfusion/embeddings.py

@@ -2,7 +2,7 @@
 Finalfusion Embeddings
 """
 from os import PathLike
-from typing import Optional, Tuple, List, Union, Any
+from typing import Optional, Tuple, List, Union, Any, Iterator
 
 import numpy as np
 
@@ -421,23 +421,24 @@ def bucket_to_explicit(self) -> 'Embeddings':
     def __contains__(self, item):
         return item in self._vocab
 
-    def __iter__(self):
+    def __iter__(self) -> Union[Iterator[Tuple[str, np.ndarray]], Iterator[
+            Tuple[str, np.ndarray, float]]]:
         if self._norms is not None:
-            return zip(self._vocab.words, self._storage, self._norms)
-        return zip(self._vocab.words, self._storage)
+            return zip(self._vocab, self._storage, self._norms)
+        return zip(self._vocab, self._storage)
 
     def _embedding(self,
                    idx: Union[int, List[int]],
                    out: Optional[np.ndarray] = None
                    ) -> Tuple[np.ndarray, Optional[float]]:
-        res = self._storage[idx]
+        res = self._storage[idx]  # type: np.ndarray
         if res.ndim == 1:
             if out is not None:
                 out[:] = res
             else:
                 out = res
             if self._norms is not None:
-                norm = self._norms[idx]
+                norm = self._norms[idx]  # type: Optional[float]
             else:
                 norm = None
         else:

src/finalfusion/metadata.py

@@ -31,7 +31,7 @@ class Metadata(dict, Chunk):
     'other value'
     """
     @staticmethod
-    def chunk_identifier():
+    def chunk_identifier() -> ChunkIdentifier:
         return ChunkIdentifier.Metadata
 
     @staticmethod

src/finalfusion/norms.py

@@ -5,21 +5,21 @@
 import struct
 from os import PathLike
 import sys
-from typing import BinaryIO, Union
+from typing import BinaryIO, Union, List, Collection
 
 import numpy as np
 
 from finalfusion.io import Chunk, ChunkIdentifier, find_chunk, TypeId, FinalfusionFormatError, \
     _pad_float32, _write_binary, _read_required_binary, _serialize_array_as_le
 
 
-class Norms(np.ndarray, Chunk):
+class Norms(np.ndarray, Chunk, Collection[float]):
     """
     Norms Chunk.
 
     Norms subclass `numpy.ndarray`, all typical numpy operations are available.
     """
-    def __new__(cls, array: np.array):
+    def __new__(cls, array: np.ndarray):
         """
         Construct new Norms.
 
@@ -46,7 +46,7 @@ def __new__(cls, array: np.array):
         return array.view(cls)
 
     @staticmethod
-    def chunk_identifier():
+    def chunk_identifier() -> ChunkIdentifier:
         return ChunkIdentifier.NdNorms
 
     @staticmethod
@@ -72,10 +72,12 @@ def write_chunk(self, file: BinaryIO):
                       int(TypeId.f32))
         _serialize_array_as_le(file, self)
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: Union[int, slice, List[int], np.ndarray]
+                    ) -> Union[float, 'Norms']:
         if isinstance(key, slice):
             return Norms(super().__getitem__(key))
-        return super().__getitem__(key)
+        norm = super().__getitem__(key)  # type: float
+        return norm
 
 
 def load_norms(file: Union[str, bytes, int, PathLike]):

src/finalfusion/scripts/convert.py

@@ -16,7 +16,7 @@
 from finalfusion.scripts.util import Format
 
 
-def main():  # pylint: disable=missing-function-docstring
+def main() -> None:  # pylint: disable=missing-function-docstring
     formats = ["word2vec", "finalfusion", "fasttext", "text", "textdims"]
     parser = argparse.ArgumentParser(prog="ffp-convert",
                                      description="Convert embeddings.")

src/finalfusion/storage/ndarray.py

@@ -5,7 +5,7 @@
 import struct
 from os import PathLike
 import sys
-from typing import BinaryIO, Tuple, Union
+from typing import BinaryIO, Tuple, Union, Iterator
 
 import numpy as np
 
@@ -45,7 +45,7 @@ def load(cls, file: BinaryIO, mmap: bool = False) -> 'NdArray':
         return cls.mmap_chunk(file) if mmap else cls.read_chunk(file)
 
     @staticmethod
-    def chunk_identifier():
+    def chunk_identifier() -> ChunkIdentifier:
         return ChunkIdentifier.NdArray
 
     @staticmethod
@@ -118,10 +118,13 @@ def write_chunk(self, file: BinaryIO):
         _write_binary(file, f"{padding}x")
         _serialize_array_as_le(file, self)
 
-    def __getitem__(self, key):
-        if isinstance(key, slice):
-            return super().__getitem__(key)
-        return super().__getitem__(key).view(np.ndarray)
+    def __getitem__(self, index) -> Union['NdArray', np.ndarray]:
+        if isinstance(index, slice):
+            return super().__getitem__(index)
+        return np.ndarray.__getitem__(self, index).view(np.ndarray)
+
+    def __iter__(self) -> Iterator[np.ndarray]:
+        return iter(self.view(np.ndarray))
 
 
 def load_ndarray(file: Union[str, bytes, int, PathLike],