Use Enum for coloring options

git_sim/__main__.py

@@ -1,13 +1,13 @@
-import pathlib
-import typer
+import datetime
 import os
+import pathlib
 import sys
-import datetime
 import time
 
-import git_sim.commands
+import typer
 
-from git_sim.settings import ImgFormat, VideoFormat, settings
+import git_sim.commands
+from git_sim.settings import ColorByOptions, ImgFormat, VideoFormat, settings
 
 app = typer.Typer(context_settings={"help_option_names": ["-h", "--help"]})
 
@@ -126,17 +126,17 @@ def main(
         settings.all,
         help="Display all local branches in the log output",
     ),
-    color_by: str = typer.Option(
+    color_by: ColorByOptions = typer.Option(
         settings.color_by,
-        help="Color commits by parameter, such as author",
+        help="Color commits by parameter",
     ),
     highlight_commit_messages: bool = typer.Option(
         settings.highlight_commit_messages,
         help="Make the displayed commit messages more prominent",
     ),
 ):
     import git
-    from manim import config, WHITE
+    from manim import WHITE, config
 
     settings.animate = animate
     settings.n = n

git_sim/enums.py

@@ -12,3 +12,20 @@ class StashSubCommand(Enum):
     POP = "pop"
     APPLY = "apply"
     PUSH = "push"
+
+
+class ColorByOptions(Enum):
+    author = "author"
+    branch = "branch"
+    notlocal1 = "notlocal1"
+    notlocal2 = "notlocal2"
+
+
+class VideoFormat(str, Enum):
+    mp4 = "mp4"
+    webm = "webm"
+
+
+class ImgFormat(str, Enum):
+    jpg = "jpg"
+    png = "png"

git_sim/git_sim_base_command.py

@@ -1,16 +1,17 @@
-import platform
-import sys
 import os
-import tempfile
+import platform
 import shutil
 import stat
+import sys
+import tempfile
 
 import git
 import manim as m
 import numpy
 from git.exc import GitCommandError, InvalidGitRepositoryError
 from git.repo import Repo
 
+from git_sim.enums import ColorByOptions
 from git_sim.settings import settings
 
 
@@ -1179,7 +1180,7 @@ def create_zone_text(
             thirdColumnFilesDict[f] = text
 
     def color_by(self, offset=0):
-        if settings.color_by == "author":
+        if settings.color_by == ColorByOptions.author:
             sorted_authors = sorted(
                 self.author_groups.keys(),
                 key=lambda k: len(self.author_groups[k]),
@@ -1208,17 +1209,17 @@ def color_by(self, offset=0):
             self.recenter_frame()
             self.scale_frame()
 
-        elif settings.color_by == "branch":
+        elif settings.color_by == ColorByOptions.branch:
             pass
 
-        elif settings.color_by == "notlocal1":
+        elif settings.color_by == ColorByOptions.notlocal1:
             for commit_id in self.drawnCommits:
                 try:
                     self.orig_repo.commit(commit_id)
                 except ValueError:
                     self.drawnCommits[commit_id].set_color(m.GOLD)
 
-        elif settings.color_by == "notlocal2":
+        elif settings.color_by == ColorByOptions.notlocal2:
             for commit_id in self.drawnCommits:
                 if not self.orig_repo.is_ancestor(commit_id, "HEAD"):
                     self.drawnCommits[commit_id].set_color(m.GOLD)

git_sim/settings.py

@@ -1,18 +1,9 @@
 import pathlib
-
-from enum import Enum
 from typing import List, Union
-from pydantic import BaseSettings
-
-
-class VideoFormat(str, Enum):
-    mp4 = "mp4"
-    webm = "webm"
 
+from pydantic import BaseSettings
 
-class ImgFormat(str, Enum):
-    jpg = "jpg"
-    png = "png"
+from git_sim.enums import ColorByOptions, ImgFormat, VideoFormat
 
 
 class Settings(BaseSettings):
@@ -46,7 +37,7 @@ class Settings(BaseSettings):
     invert_branches = False
     hide_merged_branches = False
     all = False
-    color_by: Union[str, None] = None
+    color_by: Union[ColorByOptions, None] = None
     highlight_commit_messages = False
 
     class Config:

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="git-sim",
-    version="0.2.8",
+    version="0.2.9",
     author="Jacob Stopak",
     author_email="jacob@initialcommit.io",
     description="Simulate Git commands on your own repos by generating an image (default) or video visualization depicting the command's behavior.",