Fix mypy issues

Author: SpirinEgor

code2seq/data/path_context.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import List, Iterable, Tuple, Optional
+from typing import Iterable, Tuple, Optional, Sequence
 
 import torch
 
@@ -14,12 +14,12 @@ class Path:
 @dataclass
 class LabeledPathContext:
     label: torch.Tensor  # [max label parts]
-    path_contexts: List[Path]
+    path_contexts: Sequence[Path]
 
 
 class BatchedLabeledPathContext:
-    def __init__(self, samples: List[Optional[LabeledPathContext]]):
-        samples = [s for s in samples if s is not None]
+    def __init__(self, all_samples: Sequence[Optional[LabeledPathContext]]):
+        samples = [s for s in all_samples if s is not None]
 
         # [max label parts; batch size]
         self.labels = torch.cat([s.label.unsqueeze(1) for s in samples], dim=1)
@@ -59,12 +59,13 @@ class TypedPath(Path):
 
 @dataclass
 class LabeledTypedPathContext(LabeledPathContext):
-    path_contexts: List[TypedPath]
+    path_contexts: Sequence[TypedPath]
 
 
 class BatchedLabeledTypedPathContext(BatchedLabeledPathContext):
-    def __init__(self, samples: List[Optional[LabeledTypedPathContext]]):
-        super().__init__(samples)
+    def __init__(self, all_samples: Sequence[Optional[LabeledTypedPathContext]]):
+        super().__init__(all_samples)
+        samples = [s for s in all_samples if s is not None]
         # [max type parts; n contexts]
         self.from_type = torch.cat([path.from_type.unsqueeze(1) for s in samples for path in s.path_contexts], dim=1)
         # [max type parts; n contexts]

code2seq/data/path_context_data_module.py

@@ -53,6 +53,8 @@ def collate_wrapper(batch: List[Optional[LabeledPathContext]]) -> BatchedLabeled
         return BatchedLabeledPathContext(batch)
 
     def _create_dataset(self, holdout_file: str, random_context: bool) -> PathContextDataset:
+        if self._vocabulary is None:
+            raise RuntimeError(f"Setup vocabulary before creating data loaders")
         return PathContextDataset(holdout_file, self._config, self._vocabulary, random_context)
 
     def _shared_dataloader(self, holdout: str) -> DataLoader:

code2seq/data/typed_path_context_data_module.py

@@ -17,10 +17,14 @@ def __init__(self, data_dir: str, config: DictConfig):
         super().__init__(data_dir, config)
 
     @staticmethod
-    def collate_wrapper(batch: List[Optional[LabeledTypedPathContext]]) -> BatchedLabeledTypedPathContext:
+    def collate_wrapper(  # type: ignore[override]
+        batch: List[Optional[LabeledTypedPathContext]],
+    ) -> BatchedLabeledTypedPathContext:
         return BatchedLabeledTypedPathContext(batch)
 
     def _create_dataset(self, holdout_file: str, random_context: bool) -> TypedPathContextDataset:
+        if self._vocabulary is None:
+            raise RuntimeError(f"Setup vocabulary before creating data loaders")
         return TypedPathContextDataset(holdout_file, self._config, self._vocabulary, random_context)
 
     def setup(self, stage: Optional[str] = None):

code2seq/data/typed_path_context_dataset.py

@@ -10,7 +10,7 @@
 class TypedPathContextDataset(PathContextDataset):
     def __init__(self, data_file: str, config: DictConfig, vocabulary: TypedVocabulary, random_context: bool):
         super().__init__(data_file, config, vocabulary, random_context)
-        self._vocab = vocabulary
+        self._vocab: TypedVocabulary = vocabulary
 
     def _get_path(self, raw_path: List[str]) -> TypedPath:
         return TypedPath(

code2seq/model/code2class.py

@@ -4,6 +4,7 @@
 from commode_utils.modules import Classifier
 from omegaconf import DictConfig
 from pytorch_lightning import LightningModule
+from pytorch_lightning.utilities.types import EPOCH_OUTPUT
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import _LRScheduler
 from torchmetrics import Metric, Accuracy, MetricCollection
@@ -77,18 +78,19 @@ def test_step(self, batch: BatchedLabeledPathContext, batch_idx: int) -> Dict:
 
     # ========== ON EPOCH END ==========
 
-    def _shared_epoch_end(self, outputs: List[Dict], step: str):
+    def _shared_epoch_end(self, outputs: EPOCH_OUTPUT, step: str):
+        assert isinstance(outputs, dict)
         with torch.no_grad():
             mean_loss = torch.stack([out[f"{step}/loss"] for out in outputs]).mean()
             accuracy = self.__metrics[f"{step}_acc"].compute()
             log = {f"{step}/loss": mean_loss, f"{step}/accuracy": accuracy}
         self.log_dict(log, on_step=False, on_epoch=True)
 
-    def training_epoch_end(self, outputs: List[Dict]):
+    def training_epoch_end(self, outputs: EPOCH_OUTPUT):
         self._shared_epoch_end(outputs, "train")
 
-    def validation_epoch_end(self, outputs: List[Dict]):
+    def validation_epoch_end(self, outputs: EPOCH_OUTPUT):
         self._shared_epoch_end(outputs, "val")
 
-    def test_epoch_end(self, outputs: List[Dict]):
+    def test_epoch_end(self, outputs: EPOCH_OUTPUT):
         self._shared_epoch_end(outputs, "test")

code2seq/model/modules/typed_path_encoder.py

@@ -27,7 +27,7 @@ def _calculate_concat_size(embedding_size: int, rnn_size: int, num_directions: i
     def _type_embedding(self, types: torch.Tensor) -> torch.Tensor:
         return self.type_embedding(types).sum(0)
 
-    def forward(
+    def forward(  # type: ignore
         self,
         from_type: torch.Tensor,
         from_token: torch.Tensor,

code2seq/model/typed_code2seq.py

@@ -18,7 +18,7 @@ def __init__(
         teacher_forcing: float = 0.0,
     ):
         super().__init__(model_config, optimizer_config, vocabulary, teacher_forcing)
-        self._vocabulary = vocabulary
+        self._vocabulary: TypedVocabulary = vocabulary
 
     def _get_encoder(self, config: DictConfig) -> PathEncoder:
         return TypedPathEncoder(
@@ -46,7 +46,7 @@ def forward(  # type: ignore
         output_logits = self._decoder(encoded_paths, contexts_per_label, output_length, target_sequence)
         return output_logits
 
-    def logits_from_batch(
+    def logits_from_batch(  # type: ignore[override]
         self, batch: BatchedLabeledTypedPathContext, target_sequence: Optional[torch.Tensor]
     ) -> torch.Tensor:
         return self(