Add missing type annotations

qutech · Jul 19, 2018 · d68ce6c · d68ce6c
1 parent e6d6632
commit d68ce6c
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/qctoolkit/_program/transformation.py b/qctoolkit/_program/transformation.py
@@ -1,4 +1,4 @@
-from typing import Mapping, Set
+from typing import Mapping, Set, Dict
 from abc import abstractmethod
 
 import numpy as np
@@ -44,12 +44,12 @@ def __call__(self, time: np.ndarray, data: pd.DataFrame) -> Mapping[ChannelID, n
 
         return self._matrix @ data_in
 
-    def get_output_channels(self, input_channels: Set[ChannelID]):
+    def get_output_channels(self, input_channels: Set[ChannelID]) -> Set[ChannelID]:
         if input_channels != set(self._matrix.columns):
             raise KeyError('Invalid input channels', input_channels, set(self._matrix.columns))
 
         return set(self._matrix.index)
 
     @property
-    def compare_key(self):
+    def compare_key(self) -> Dict[ChannelID, Dict[ChannelID, float]]:
         return self._matrix.to_dict()