How to Roll Back to a Previous Image Tag in Python

"""Demonstrates a rollback pattern for a Docker-style image tag registry."""

from dataclasses import dataclass, field


@dataclass
class ImageRegistry:
    """A minimal mock registry tracking current tags per image."""

    tags: dict[str, list[str]] = field(default_factory=dict)

    def push(self, image: str, tag: str) -> None:
        """Add a tag to the image's history."""
        self.tags.setdefault(image, []).append(tag)

    def rollback(self, image: str) -> str | None:
        """Roll back to the previous tag (removes current, returns new current)."""
        history = self.tags.get(image, [])
        if len(history) <= 1:
            return None
        history.pop()
        return history[-1]

    def current(self, image: str) -> str | None:
        """Return the latest tag for an image, if any."""
        return self.tags.get(image, [None])[-1]


if __name__ == "__main__":
    registry = ImageRegistry()
    registry.push("web/api", "v1.0.0")
    registry.push("web/api", "v1.0.1")
    registry.push("web/api", "v1.0.2")

    print("Current before rollback:", registry.current("web/api"))
    print("New current after rollback:", registry.rollback("web/api"))
    print("Rollback result:", registry.rollback("web/api"))
    print("Current after second rollback:", registry.current("web/api"))
    print("Rollback on single tag returns:", registry.rollback("web/api"))
    print("Current at end:", registry.current("web/api"))