♻️ Refactor to use enum

Author: tiangolo

scripts/prepare_release.py

@@ -2,16 +2,23 @@
 
 import re
 from datetime import date
+from enum import Enum
 from pathlib import Path
-from typing import Annotated, Literal
+from typing import Annotated
 
 import typer
 
 VERSION_PATTERN = re.compile(r'(?m)^__version__ = "(\d+\.\d+\.\d+)"$')
 VERSION_HEADING_PATTERN = re.compile(r"(?m)^## (\d+\.\d+\.\d+)(?: \([^)]+\))?$")
 RELEASE_NOTES_HEADER = "# Release Notes\n\n"
 LATEST_CHANGES_HEADER = "## Latest Changes"
-BumpType = Literal["major", "minor", "patch"]
+
+
+class BumpType(str, Enum):
+    major = "major"
+    minor = "minor"
+    patch = "patch"
+
 
 app = typer.Typer()
 
@@ -36,9 +43,9 @@ def get_current_version(content: str, version_file: Path) -> str:
 
 def bump_version(version: str, bump: BumpType) -> str:
     major, minor, patch = parse_version(version)
-    if bump == "major":
+    if bump is BumpType.major:
         return f"{major + 1}.0.0"
-    if bump == "minor":
+    if bump is BumpType.minor:
         return f"{major}.{minor + 1}.0"
     return f"{major}.{minor}.{patch + 1}"
 

tests/test_prepare_release.py

@@ -19,9 +19,9 @@
 @pytest.mark.parametrize(
     ("current_version", "bump", "new_version"),
     [
-        ("0.0.24", "major", "1.0.0"),
-        ("0.0.24", "minor", "0.1.0"),
-        ("0.0.24", "patch", "0.0.25"),
+        ("0.0.24", BumpType.major, "1.0.0"),
+        ("0.0.24", BumpType.minor, "0.1.0"),
+        ("0.0.24", BumpType.patch, "0.0.25"),
     ],
 )
 def test_bump_version(current_version: str, bump: BumpType, new_version: str) -> None: