From 148331f04b2479970c42c94ca0c83fe88962bbbb Mon Sep 17 00:00:00 2001
From: Acme Contributor <noreply@google.com>
Date: Wed, 3 Jan 2024 08:44:18 -0800
Subject: [PATCH] Added pytype None checks to savers.py

PiperOrigin-RevId: 595412439
Change-Id: I1367163d0c2cd08c22446bc9e439b946628d4e4c
---
 acme/tf/savers.py | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/acme/tf/savers.py b/acme/tf/savers.py
index f3b623e71d..1b33478afd 100644
--- a/acme/tf/savers.py
+++ b/acme/tf/savers.py
@@ -150,23 +150,38 @@ def save(self, force: bool = False) -> bool:
         time.time() - self._last_saved < 60 * self._time_delta_minutes):
       return False
 
+    checkpoint_manager: tf.train.CheckpointManager = self.checkpoint_manager
     # Save any checkpoints.
-    logging.info('Saving checkpoint: %s', self._checkpoint_manager.directory)
-    self._checkpoint_manager.save()
+    logging.info('Saving checkpoint: %s', checkpoint_manager.directory)
+    checkpoint_manager.save()
     self._last_saved = time.time()
 
     return True
 
   def restore(self):
+    """Restore from most recent checkpoint."""
+
     # Restore from the most recent checkpoint (if it exists).
-    checkpoint_to_restore = self._checkpoint_manager.latest_checkpoint
+    checkpoint_to_restore = self.checkpoint_manager.latest_checkpoint
     logging.info('Attempting to restore checkpoint: %s',
                  checkpoint_to_restore)
     self._checkpoint.restore(checkpoint_to_restore)
 
   @property
   def directory(self):
-    return self._checkpoint_manager.directory
+    return self.checkpoint_manager.directory
+
+  @property
+  def checkpoint_manager(self) -> tf.train.CheckpointManager:
+    if not self._enable_checkpointing:
+      raise ValueError(
+          'Check-point not enabled. No checkpoint manager available.'
+      )
+
+    # At this point, _enable_checkpointing is true, so _checkpoint_manager
+    # should not be None.
+    assert self._checkpoint_manager is not None
+    return self._checkpoint_manager
 
 
 class CheckpointingRunner(core.Worker):
@@ -332,6 +347,8 @@ def __init__(self):
 
   @tf.function
   def __call__(self, *args, **kwargs):
+    if self._module is None:
+      raise ValueError('_module not set')
     return self._module(*args, **kwargs)
 
   @property