Added type hints

This commit is contained in:
Andrew Murray 2024-06-08 22:26:28 +10:00
parent 5e0fe8f1b8
commit 14a32650dd

View File

@ -25,7 +25,7 @@ from . import Image
@lru_cache @lru_cache
def getrgb(color): def getrgb(color: str) -> tuple[int, int, int] | tuple[int, int, int, int]:
""" """
Convert a color string to an RGB or RGBA tuple. If the string cannot be Convert a color string to an RGB or RGBA tuple. If the string cannot be
parsed, this function raises a :py:exc:`ValueError` exception. parsed, this function raises a :py:exc:`ValueError` exception.
@ -44,8 +44,10 @@ def getrgb(color):
if rgb: if rgb:
if isinstance(rgb, tuple): if isinstance(rgb, tuple):
return rgb return rgb
colormap[color] = rgb = getrgb(rgb) rgb_tuple = getrgb(rgb)
return rgb assert len(rgb_tuple) == 3
colormap[color] = rgb_tuple
return rgb_tuple
# check for known string formats # check for known string formats
if re.match("#[a-f0-9]{3}$", color): if re.match("#[a-f0-9]{3}$", color):
@ -88,15 +90,15 @@ def getrgb(color):
if m: if m:
from colorsys import hls_to_rgb from colorsys import hls_to_rgb
rgb = hls_to_rgb( rgb_floats = hls_to_rgb(
float(m.group(1)) / 360.0, float(m.group(1)) / 360.0,
float(m.group(3)) / 100.0, float(m.group(3)) / 100.0,
float(m.group(2)) / 100.0, float(m.group(2)) / 100.0,
) )
return ( return (
int(rgb[0] * 255 + 0.5), int(rgb_floats[0] * 255 + 0.5),
int(rgb[1] * 255 + 0.5), int(rgb_floats[1] * 255 + 0.5),
int(rgb[2] * 255 + 0.5), int(rgb_floats[2] * 255 + 0.5),
) )
m = re.match( m = re.match(
@ -105,15 +107,15 @@ def getrgb(color):
if m: if m:
from colorsys import hsv_to_rgb from colorsys import hsv_to_rgb
rgb = hsv_to_rgb( rgb_floats = hsv_to_rgb(
float(m.group(1)) / 360.0, float(m.group(1)) / 360.0,
float(m.group(2)) / 100.0, float(m.group(2)) / 100.0,
float(m.group(3)) / 100.0, float(m.group(3)) / 100.0,
) )
return ( return (
int(rgb[0] * 255 + 0.5), int(rgb_floats[0] * 255 + 0.5),
int(rgb[1] * 255 + 0.5), int(rgb_floats[1] * 255 + 0.5),
int(rgb[2] * 255 + 0.5), int(rgb_floats[2] * 255 + 0.5),
) )
m = re.match(r"rgba\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\)$", color) m = re.match(r"rgba\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\)$", color)
@ -124,7 +126,7 @@ def getrgb(color):
@lru_cache @lru_cache
def getcolor(color, mode: str) -> tuple[int, ...]: def getcolor(color: str, mode: str) -> int | tuple[int, ...]:
""" """
Same as :py:func:`~PIL.ImageColor.getrgb` for most modes. However, if Same as :py:func:`~PIL.ImageColor.getrgb` for most modes. However, if
``mode`` is HSV, converts the RGB value to a HSV value, or if ``mode`` is ``mode`` is HSV, converts the RGB value to a HSV value, or if ``mode`` is
@ -136,33 +138,34 @@ def getcolor(color, mode: str) -> tuple[int, ...]:
:param color: A color string :param color: A color string
:param mode: Convert result to this mode :param mode: Convert result to this mode
:return: ``(graylevel[, alpha]) or (red, green, blue[, alpha])`` :return: ``graylevel, (graylevel, alpha) or (red, green, blue[, alpha])``
""" """
# same as getrgb, but converts the result to the given mode # same as getrgb, but converts the result to the given mode
color, alpha = getrgb(color), 255 rgb, alpha = getrgb(color), 255
if len(color) == 4: if len(rgb) == 4:
color, alpha = color[:3], color[3] alpha = rgb[3]
rgb = rgb[:3]
if mode == "HSV": if mode == "HSV":
from colorsys import rgb_to_hsv from colorsys import rgb_to_hsv
r, g, b = color r, g, b = rgb
h, s, v = rgb_to_hsv(r / 255, g / 255, b / 255) h, s, v = rgb_to_hsv(r / 255, g / 255, b / 255)
return int(h * 255), int(s * 255), int(v * 255) return int(h * 255), int(s * 255), int(v * 255)
elif Image.getmodebase(mode) == "L": elif Image.getmodebase(mode) == "L":
r, g, b = color r, g, b = rgb
# ITU-R Recommendation 601-2 for nonlinear RGB # ITU-R Recommendation 601-2 for nonlinear RGB
# scaled to 24 bits to match the convert's implementation. # scaled to 24 bits to match the convert's implementation.
color = (r * 19595 + g * 38470 + b * 7471 + 0x8000) >> 16 graylevel = (r * 19595 + g * 38470 + b * 7471 + 0x8000) >> 16
if mode[-1] == "A": if mode[-1] == "A":
return color, alpha return graylevel, alpha
else: return graylevel
if mode[-1] == "A": elif mode[-1] == "A":
return color + (alpha,) return rgb + (alpha,)
return color return rgb
colormap = { colormap: dict[str, str | tuple[int, int, int]] = {
# X11 colour table from https://drafts.csswg.org/css-color-4/, with # X11 colour table from https://drafts.csswg.org/css-color-4/, with
# gray/grey spelling issues fixed. This is a superset of HTML 4.0 # gray/grey spelling issues fixed. This is a superset of HTML 4.0
# colour names used in CSS 1. # colour names used in CSS 1.