memo: Add typing to some callback classes
@@ -24,6 +24,9 @@ import os
import re
import sys
import time
+from typing import Iterable
+from typing import Optional
+from typing import Union
import numpy as np
import tensorflow.compat.v2 as tf
@@ -934,7 +937,7 @@ class BaseLogger(Callback):
All others will be averaged in `on_epoch_end`.
"""
- def __init__(self, stateful_metrics=None):
+ def __init__(self, stateful_metrics: Optional[Iterable[str]] = None):
super().__init__()
self.stateful_metrics = set(stateful_metrics or [])
@@ -1009,7 +1012,13 @@ class ProgbarLogger(Callback):
ValueError: In case of invalid `count_mode`.
"""
- def __init__(self, count_mode="samples", stateful_metrics=None):
+ def __init__(
+ self,
+ count_mode: str = "samples",
+ stateful_metrics: Optional[Iterable[str]] = None,
+ ):
+ # when we drop support for python 3.7, replace 'count_mode: str'
+ # with 'count_mode: Literal["samples", "steps"]'
super().__init__()
self._supports_tf_logs = True
if count_mode == "samples":
@@ -1318,15 +1327,17 @@ class ModelCheckpoint(Callback):
def __init__(
self,
- filepath,
- monitor="val_loss",
- verbose=0,
- save_best_only=False,
- save_weights_only=False,
- mode="auto",
- save_freq="epoch",
- options=None,
- initial_value_threshold=None,
+ filepath: Union[str, os.PathLike],
+ monitor: str = "val_loss",
+ verbose: int = 0,
+ save_best_only: bool = False,
+ save_weights_only: bool = False,
+ mode: str = "auto",
+ save_freq: Union[int, str] = "epoch",
+ options: Union[
+ tf.train.CheckpointOptions, tf.saved_model.SaveOptions, None
+ ] = None,
+ initial_value_threshold: Optional[float] = None,
**kwargs,
):
super().__init__()
