From dd474abfe878509f414012ad09e12fdb4afcb292 Mon Sep 17 00:00:00 2001
From: Yixing Lao <yixing.lao@gmail.com>
Date: Fri, 10 Jan 2025 20:49:58 +0800
Subject: [PATCH] refactor: add type hints for literal strings (#84)

---
 camtools/image.py  |  4 ++--
 camtools/io.py     |  4 ++--
 camtools/metric.py |  4 ++--
 camtools/render.py | 12 ++++++------
 4 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/camtools/image.py b/camtools/image.py
index 3e79c470..a715779e 100644
--- a/camtools/image.py
+++ b/camtools/image.py
@@ -6,7 +6,7 @@
 import cv2
 from . import sanity
 from . import colormap
-from typing import Tuple, List, Optional, Union
+from typing import Tuple, List, Optional, Union, Literal
 from jaxtyping import Float, UInt8, UInt16, Int
 
 
@@ -778,7 +778,7 @@ def get_scales(im_height, max_lines, font, line_text_h_ratio):
 
 def vstack_images(
     ims: List[Float[np.ndarray, "h w 3"]],
-    alignment: str = "left",
+    alignment: Literal["left", "center", "right"] = "left",
     background_color: Tuple[float, float, float] = (1.0, 1.0, 1.0),
 ) -> Float[np.ndarray, "h_stacked w_stacked 3"]:
     """
diff --git a/camtools/io.py b/camtools/io.py
index 46a52263..722f0154 100644
--- a/camtools/io.py
+++ b/camtools/io.py
@@ -5,7 +5,7 @@
 import cv2
 import numpy as np
 from pathlib import Path
-from typing import Union, Optional
+from typing import Union, Optional, Literal
 from jaxtyping import UInt8, Float
 
 from . import sanity
@@ -185,7 +185,7 @@ def imwrite_depth(
 
 def imread(
     im_path: Union[str, Path],
-    alpha_mode: Optional[str] = None,
+    alpha_mode: Optional[Literal["keep", "ignore", "white", "black"]] = None,
 ) -> Union[
     Float[np.ndarray, "h w"],
     Float[np.ndarray, "h w 3"],
diff --git a/camtools/metric.py b/camtools/metric.py
index 86febfea..0bbf1784 100644
--- a/camtools/metric.py
+++ b/camtools/metric.py
@@ -6,7 +6,7 @@
 from skimage.metrics import peak_signal_noise_ratio
 from skimage.metrics import structural_similarity
 from pathlib import Path
-from typing import Tuple, Optional, Union
+from typing import Tuple, Optional, Union, Literal
 from jaxtyping import Float
 
 from . import image
@@ -200,7 +200,7 @@ def load_im_pd_im_gt_im_mask_for_eval(
     im_pd_path: Union[str, Path],
     im_gt_path: Union[str, Path],
     im_mask_path: Optional[Union[str, Path]] = None,
-    alpha_mode: str = "white",
+    alpha_mode: Literal["white", "keep"] = "white",
 ) -> Tuple[
     Float[np.ndarray, "h w 3"], Float[np.ndarray, "h w 3"], Float[np.ndarray, "h w"]
 ]:
diff --git a/camtools/render.py b/camtools/render.py
index a8960b78..c88564d3 100644
--- a/camtools/render.py
+++ b/camtools/render.py
@@ -2,7 +2,7 @@
 Functions for controlled rendering of 3D geometries to images or depth images.
 """
 
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Literal
 
 import numpy as np
 import open3d as o3d
@@ -422,7 +422,7 @@ class _TextRenderer:
         "mono": None,
     }
 
-    def __init__(self, font_type: str = "tex"):
+    def __init__(self, font_type: Literal["tex", "serif", "sans", "mono"] = "tex"):
         """
         Initializes the renderer with a specific font type.
         """
@@ -543,10 +543,10 @@ def render(
 def render_text(
     text: str,
     font_size: int = 72,
-    font_type: str = "tex",
+    font_type: Literal["tex", "serif", "sans", "mono"] = "tex",
     font_color: Tuple[float, float, float] = (0, 0, 0),
     tight_layout: bool = False,
-    multiline_alignment: str = "left",
+    multiline_alignment: Literal["left", "center", "right"] = "left",
     padding_tblr: Tuple[int, int, int, int] = (0, 0, 0, 0),
 ) -> Float[np.ndarray, "h w"]:
     """
@@ -604,9 +604,9 @@ def render_text(
 def render_texts(
     texts: List[str],
     font_size: int = 72,
-    font_type: str = "tex",
+    font_type: Literal["tex", "serif", "sans", "mono"] = "tex",
     font_color: Tuple[float, float, float] = (0.0, 0.0, 0.0),
-    multiline_alignment: str = "center",
+    multiline_alignment: Literal["left", "center", "right"] = "center",
     same_height: bool = False,
     same_width: bool = False,
     padding_tblr: Tuple[int, int, int, int] = (0, 0, 0, 0),