TYP: Add overloads for MathTextParser return types

QuLogic · QuLogic · commit 609ca9768221 · 2023-08-30T03:36:46.000-04:00
diff --git a/lib/matplotlib/backends/backend_agg.py b/lib/matplotlib/backends/backend_agg.py
@@ -71,7 +71,7 @@ def __init__(self, width, height, dpi):
         self._filter_renderers = []
 
         self._update_methods()
-        self.mathtext_parser = MathTextParser('Agg')
+        self.mathtext_parser = MathTextParser('agg')
 
         self.bbox = Bbox.from_bounds(0, 0, self.width, self.height)
 
diff --git a/lib/matplotlib/mathtext.pyi b/lib/matplotlib/mathtext.pyi
@@ -1,5 +1,8 @@
 import os
+from typing import Generic, IO, Literal, TypeVar, overload
+
 from matplotlib.font_manager import FontProperties
+from matplotlib.typing import ColorType
 
 # Re-exported API from _mathtext.
 from ._mathtext import (
@@ -8,14 +11,16 @@ from ._mathtext import (
     get_unicode_index as get_unicode_index,
 )
 
-from typing import IO, Literal
-from matplotlib.typing import ColorType
+_ParseType = TypeVar("_ParseType", RasterParse, VectorParse)
 
-class MathTextParser:
-    def __init__(self, output: Literal["path", "agg", "raster", "macosx"]) -> None: ...
+class MathTextParser(Generic[_ParseType]):
+    @overload
+    def __init__(self: MathTextParser[VectorParse], output: Literal["path"]) -> None: ...
+    @overload
+    def __init__(self: MathTextParser[RasterParse], output: Literal["agg", "raster", "macosx"]) -> None: ...
     def parse(
         self, s: str, dpi: float = ..., prop: FontProperties | None = ..., *, antialiased: bool | None = ...
-    ) -> RasterParse | VectorParse: ...
+    ) -> _ParseType: ...
 
 def math_to_image(
     s: str,
diff --git a/lib/matplotlib/textpath.pyi b/lib/matplotlib/textpath.pyi
@@ -1,6 +1,6 @@
 from matplotlib.font_manager import FontProperties
 from matplotlib.ft2font import FT2Font
-from matplotlib.mathtext import MathTextParser
+from matplotlib.mathtext import MathTextParser, VectorParse
 from matplotlib.path import Path
 
 import numpy as np
@@ -10,7 +10,7 @@ from typing import Literal
 class TextToPath:
     FONT_SCALE: float
     DPI: float
-    mathtext_parser: MathTextParser
+    mathtext_parser: MathTextParser[VectorParse]
     def __init__(self) -> None: ...
     def get_text_width_height_descent(
         self, s: str, prop: FontProperties, ismath: bool | Literal["TeX"]
