more typing work

kittens/diff/highlight.py

@@ -5,7 +5,7 @@
 import concurrent
 import os
 import re
-from typing import List, Optional
+from typing import IO, Dict, Iterable, List, Optional, Tuple, Union, cast
 
 from pygments import highlight  # type: ignore
 from pygments.formatter import Formatter  # type: ignore
@@ -14,7 +14,7 @@ from pygments.util import ClassNotFound  # type: ignore
 
 from kitty.rgb import color_as_sgr, parse_sharp
 
-from .collect import Segment, data_for_path, lines_for_path
+from .collect import Collection, Segment, data_for_path, lines_for_path
 
 
 class StyleNotFound(Exception):
@@ -23,7 +23,7 @@ class StyleNotFound(Exception):
 
 class DiffFormatter(Formatter):
 
-    def __init__(self, style='default'):
+    def __init__(self, style: str = 'default') -> None:
         try:
             Formatter.__init__(self, style=style)
             initialized = True
@@ -32,26 +32,26 @@ class DiffFormatter(Formatter):
         if not initialized:
             raise StyleNotFound('pygments style "{}" not found'.format(style))
 
-        self.styles = {}
-        for token, style in self.style:
+        self.styles: Dict[str, Tuple[str, str]] = {}
+        for token, token_style in self.style:
             start = []
             end = []
             fstart = fend = ''
             # a style item is a tuple in the following form:
             # colors are readily specified in hex: 'RRGGBB'
-            col = style['color']
+            col = token_style['color']
             if col:
                 pc = parse_sharp(col)
                 if pc is not None:
                     start.append('38' + color_as_sgr(pc))
                     end.append('39')
-            if style['bold']:
+            if token_style['bold']:
                 start.append('1')
                 end.append('22')
-            if style['italic']:
+            if token_style['italic']:
                 start.append('3')
                 end.append('23')
-            if style['underline']:
+            if token_style['underline']:
                 start.append('4')
                 end.append('24')
             if start:
@@ -59,7 +59,7 @@ class DiffFormatter(Formatter):
                 fend = '\033[{}m'.format(';'.join(end))
             self.styles[token] = fstart, fend
 
-    def format(self, tokensource, outfile):
+    def format(self, tokensource: Iterable[Tuple[str, str]], outfile: IO[str]) -> None:
         for ttype, value in tokensource:
             not_found = True
             if value.rstrip('\n'):
@@ -84,12 +84,12 @@ class DiffFormatter(Formatter):
 formatter: Optional[DiffFormatter] = None
 
 
-def initialize_highlighter(style='default'):
+def initialize_highlighter(style: str = 'default') -> None:
     global formatter
     formatter = DiffFormatter(style)
 
 
-def highlight_data(code, filename, aliases=None):
+def highlight_data(code: str, filename: str, aliases: Optional[Dict[str, str]] = None) -> Optional[str]:
     if aliases:
         base, ext = os.path.splitext(filename)
         alias = aliases.get(ext[1:])
@@ -98,15 +98,14 @@ def highlight_data(code, filename, aliases=None):
     try:
         lexer = get_lexer_for_filename(filename, stripnl=False)
     except ClassNotFound:
-        pass
-    else:
-        return highlight(code, lexer, formatter)
+        return None
+    return cast(str, highlight(code, lexer, formatter))
 
 
 split_pat = re.compile(r'(\033\[.*?m)')
 
 
-def highlight_line(line):
+def highlight_line(line: str) -> List[Segment]:
     ans: List[Segment] = []
     current: Optional[Segment] = None
     pos = 0
@@ -124,8 +123,11 @@ def highlight_line(line):
     return ans
 
 
-def highlight_for_diff(path, aliases):
-    ans = []
+DiffHighlight = List[List[Segment]]
+
+
+def highlight_for_diff(path: str, aliases: Dict[str, str]) -> DiffHighlight:
+    ans: DiffHighlight = []
     lines = lines_for_path(path)
     hd = highlight_data('\n'.join(lines), path, aliases)
     if hd is not None:
@@ -134,9 +136,9 @@ def highlight_for_diff(path, aliases):
     return ans
 
 
-def highlight_collection(collection, aliases=None):
+def highlight_collection(collection: Collection, aliases: Optional[Dict[str, str]] = None) -> Union[str, Dict[str, DiffHighlight]]:
     jobs = {}
-    ans = {}
+    ans: Dict[str, DiffHighlight] = {}
     with concurrent.futures.ProcessPoolExecutor(max_workers=os.cpu_count()) as executor:
         for path, item_type, other_path in collection:
             if item_type != 'rename':
@@ -155,7 +157,7 @@ def highlight_collection(collection, aliases=None):
     return ans
 
 
-def main():
+def main() -> None:
     from .config import defaults
     # kitty +runpy "from kittens.diff.highlight import main; main()" file
     import sys