make mypy happy

ArneBinder · Nov 7, 2023 · 8e2d845 · 8e2d845
1 parent 74556fe
commit 8e2d845
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/pie_datasets/statistics.py b/src/pie_datasets/statistics.py
@@ -1,6 +1,6 @@
 import logging
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, Type, Union
+from typing import Any, Callable, Dict, List, Optional, Type, Union
 
 from pytorch_ie.annotations import Span
 from pytorch_ie.core import Document, DocumentStatistic
@@ -110,7 +110,7 @@ def __init__(
                 f"this collector, but be aware that the results may be wrong for your own aggregation "
                 f"functions that rely on zero values."
             )
-            self.aggregation_functions = {
+            self.aggregation_functions: Dict[str, Callable[[List], Any]] = {
                 name: func
                 for name, func in self.aggregation_functions.items()
                 if name not in ["mean", "std", "min"]
@@ -226,7 +226,7 @@ def __init__(
                 f"this collector, but be aware that the results may be wrong for your own aggregation "
                 f"functions that rely on zero values."
             )
-            self.aggregation_functions = {
+            self.aggregation_functions: Dict[str, Callable[[List], Any]] = {
                 name: func
                 for name, func in self.aggregation_functions.items()
                 if name not in ["mean", "std", "min"]