Fix typechecking

Author: Cimbali

src/prompt_toolkit/layout/containers.py

@@ -1944,7 +1944,7 @@ def render_margin(m: Margin, width: int) -> UIContent:
     def _whitespace_wrap_finder(
         self,
         ui_content: UIContent,
-        sep: str | re.Pattern = r"\s",
+        sep: str | re.Pattern[str] = r"\s",
         split: str = "remove",
         continuation: StyleAndTextTuples = [],
     ) -> WrapFinderCallable:
@@ -2024,15 +2024,22 @@ def _copy_body(
         # Maps (row, col) from the input to (y, x) screen coordinates.
         rowcol_to_yx: dict[tuple[int, int], tuple[int, int]] = {}
 
-        def find_next_wrap(remaining_width, is_input, lineno, fragment=0, char_pos=0):
+        def find_next_wrap(
+            remaining_width: int,
+            is_input: bool,
+            lineno: int,
+            fragment: int = 0,
+            char_pos: int = 0,
+        ) -> tuple[int, int, AnyFormattedText]:
             if not wrap_lines:
                 return sys.maxsize, 0, []
 
             line = ui_content.get_line(lineno)
             style0, text0, *more = line[fragment]
-            char_pos - fragment_list_len(line[:fragment])
-            line_part = [(style0, text0[char_pos:], *more), *line[fragment + 1 :]]
-            line_width = [fragment_list_width([fragment]) for fragment in line_part]
+            char_pos -= fragment_list_len(line[:fragment])
+            line_part = [(style0, text0[char_pos:]), *line[fragment + 1 :]]
+            line_width = [fragment_list_width([frag]) for frag in line_part]
+            line_width = [fragment_list_width([frag]) for frag in line_part]
 
             if sum(line_width) <= remaining_width:
                 return sys.maxsize, 0, []
@@ -2054,10 +2061,15 @@ def find_next_wrap(remaining_width, is_input, lineno, fragment=0, char_pos=0):
                 remaining_width -= char_width
                 max_wrap_pos += 1
 
-            if is_input and wrap_finder:
-                return wrap_finder(lineno, min_wrap_pos, max_wrap_pos)
-            else:
-                return max_wrap_pos, 0, []
+            return (
+                wrap_finder(lineno, min_wrap_pos, max_wrap_pos)
+                if is_input and wrap_finder
+                else None
+            ) or (
+                max_wrap_pos,
+                0,
+                [],
+            )
 
         def copy_line(
             line: StyleAndTextTuples,

src/prompt_toolkit/layout/menus.py

@@ -27,7 +27,7 @@
 from prompt_toolkit.utils import get_cwidth
 
 from .containers import ConditionalContainer, HSplit, ScrollOffsets, Window
-from .controls import GetLinePrefixCallable, UIContent, UIControl
+from .controls import GetLinePrefixCallable, UIContent, UIControl, WrapFinderCallable
 from .dimension import Dimension
 from .margins import ScrollbarMargin
 
@@ -80,6 +80,7 @@ def preferred_height(
         max_available_height: int,
         wrap_lines: bool,
         get_line_prefix: GetLinePrefixCallable | None,
+        wrap_finder: WrapFinderCallable | None,
     ) -> int | None:
         complete_state = get_app().current_buffer.complete_state
         if complete_state:
@@ -378,6 +379,7 @@ def preferred_height(
         max_available_height: int,
         wrap_lines: bool,
         get_line_prefix: GetLinePrefixCallable | None,
+        wrap_finder: WrapFinderCallable | None,
     ) -> int | None:
         """
         Preferred height: as much as needed in order to display all the completions.
@@ -718,6 +720,7 @@ def preferred_height(
         max_available_height: int,
         wrap_lines: bool,
         get_line_prefix: GetLinePrefixCallable | None,
+        wrap_finder: WrapFinderCallable | None,
     ) -> int | None:
         return 1