[benchmarks] typing

Author: jneuendorf

avalanche/benchmarks/utils/classification_dataset.py

@@ -16,7 +16,6 @@
 labels automatically. Concatenation and subsampling operations are optimized
 to be used frequently, as is common in replay strategies.
 """
-
 from functools import partial
 from typing import (
     List,
@@ -29,7 +28,7 @@
     Dict,
     Tuple,
     Mapping,
-    overload,
+    overload, Self,
 )
 
 import torch
@@ -64,11 +63,11 @@
 )
 
 T_co = TypeVar("T_co", covariant=True)
-TAvalancheDataset = TypeVar("TAvalancheDataset", bound="AvalancheDataset")
+TAvalancheDataset = TypeVar("TAvalancheDataset", bound=AvalancheDataset)
 TTargetType = int
 
 TClassificationDataset = TypeVar(
-    "TClassificationDataset", bound="ClassificationDataset"
+    "TClassificationDataset", bound=IDatasetWithTargets
 )
 
 
@@ -114,7 +113,7 @@ def task_pattern_indices(self) -> Dict[int, Sequence[int]]:
         return self.targets_task_labels.val_to_idx  # type: ignore
 
     @property
-    def task_set(self: TClassificationDataset) -> TaskSet[TClassificationDataset]:
+    def task_set(self) -> TaskSet[Self]:
         """Returns the dataset's ``TaskSet``, which is a mapping <task-id,
         task-dataset>."""
         return TaskSet(self)

avalanche/benchmarks/utils/utils.py

@@ -653,13 +653,15 @@ class TaskSet(Mapping[int, TAvalancheDataset], Generic[TAvalancheDataset]):
 
     """
 
+    data: TAvalancheDataset
+
     def __init__(self, data: TAvalancheDataset):
         """Constructor.
 
         :param data: original data
         """
         super().__init__()
-        self.data: TAvalancheDataset = data
+        self.data = data
 
     def __iter__(self) -> Iterator[int]:
         t_labels = self._get_task_labels_field()