import datetime import pathlib from typing import Union from anaxi.file_persistence import read_yaml, write_yaml class CheckpointManager: highest_synced_timestamp_field_key = "highest_synced_timestamp" def __init__( self, file_path: pathlib.Path, highest_synced_timestamp: Union[None, datetime.datetime], ): self.file_path = file_path self.highest_synced_timestamp = highest_synced_timestamp self.goal = None @classmethod def load_from_yml(cls, file_path: pathlib.Path) -> "CheckpointManager": contents = read_yaml(file_path) if CheckpointManager.highest_synced_timestamp_field_key not in contents.keys(): raise ValueError(f"Invalid checkpoint state contents in file: {file_path}.") return CheckpointManager( file_path=file_path, highest_synced_timestamp=datetime.datetime.fromisoformat( contents[CheckpointManager.highest_synced_timestamp_field_key] ), ) def update_yml(self): write_yaml( path=self.file_path, data={ CheckpointManager.highest_synced_timestamp_field_key: self.highest_synced_timestamp.isoformat() }, ) def set_new_goal(self, goal_timestamp: datetime.datetime) -> None: self.goal = goal_timestamp def commit_goal(self): self.highest_synced_timestamp = self.goal self.update_yml()