From c5488df67abf34a7ff28aed0145e59c7b84abed0 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Fri, 27 Feb 2026 09:44:01 -0600
Subject: [PATCH] Preserve type when applying optimize_mapper

---
 pymbolic/mapper/optimize.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/pymbolic/mapper/optimize.py b/pymbolic/mapper/optimize.py
index 9795bf5..7997bb7 100644
--- a/pymbolic/mapper/optimize.py
+++ b/pymbolic/mapper/optimize.py
@@ -262,18 +262,21 @@ def _set_and_return(
     return value
 
 
+MapperT = TypeVar("MapperT", bound=type)
+
+
 def optimize_mapper(
         *, drop_args: bool = False, drop_kwargs: bool = False,
         inline_rec: bool = False, inline_cache: bool = False,
         inline_get_cache_key: bool = False,
-        print_modified_code_file: TextIO | None = None) -> Callable[[type], type]:
+        print_modified_code_file: TextIO | None = None) -> Callable[[MapperT], MapperT]:
     """
     :param print_modified_code_file: a file-like object to which the modified
         code will be printed, or ``None``.
     """
     # This is a crime, an abomination. But a somewhat effective one.
 
-    def wrapper(cls: type) -> type:
+    def wrapper(cls: MapperT) -> MapperT:
         try:
             # Introduced in Py3.9
             ast.unparse  # noqa: B018
@@ -398,7 +401,7 @@ def wrapper(cls: type) -> type:
             "exec"),
              compile_dict)
 
-        return cast("type", compile_dict[cls.__name__])
+        return cast("MapperT", compile_dict[cls.__name__])
 
     return wrapper