more typing work

kittens/diff/search.py

@@ -3,11 +3,17 @@
 # License: GPL v3 Copyright: 2018, Kovid Goyal <kovid at kovidgoyal.net>
 
 import re
+from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Tuple
 
 from kitty.fast_data_types import wcswidth
+from kitty.options_stub import DiffOptions
 
 from ..tui.operations import styled
 
+if TYPE_CHECKING:
+    from .render import Line
+    Line
+
 
 class BadRegex(ValueError):
     pass
@@ -15,8 +21,8 @@ class BadRegex(ValueError):
 
 class Search:
 
-    def __init__(self, opts, query, is_regex, is_backward):
-        self.matches = {}
+    def __init__(self, opts: DiffOptions, query: str, is_regex: bool, is_backward: bool):
+        self.matches: Dict[int, List[Tuple[int, str]]] = {}
         self.count = 0
         self.style = styled('|', fg=opts.search_fg, bg=opts.search_bg).split('|', 1)[0]
         if not is_regex:
@@ -26,7 +32,7 @@ class Search:
         except Exception:
             raise BadRegex('Not a valid regex: {}'.format(query))
 
-    def __call__(self, diff_lines, margin_size, cols):
+    def __call__(self, diff_lines: Iterable['Line'], margin_size: int, cols: int) -> bool:
         self.matches = {}
         self.count = 0
         half_width = cols // 2
@@ -38,7 +44,7 @@ class Search:
             left, right = text[margin_size:half_width + 1], text[right_offset:]
             matches = []
 
-            def add(which, offset):
+            def add(which: str, offset: int) -> None:
                 for m in find(which):
                     before = which[:m.start()]
                     matches.append((wcswidth(before) + offset, m.group()))
@@ -50,13 +56,13 @@ class Search:
                 self.matches[i] = matches
         return bool(self.matches)
 
-    def __contains__(self, i):
+    def __contains__(self, i: int) -> bool:
         return i in self.matches
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self.count
 
-    def highlight_line(self, write, line_num):
+    def highlight_line(self, write: Callable[[str], None], line_num: int) -> bool:
         highlights = self.matches.get(line_num)
         if not highlights:
             return False