Pygments: make Formatter generic and improve format/highlight (#6819)

Co-authored-by: Akuli <[email protected]>

Author: Dreamsorcerer

stubs/Pygments/pygments/__init__.pyi

@@ -1,5 +1,16 @@
-from typing import Any
+from _typeshed import SupportsWrite
+from typing import TypeVar, overload
+
+from pygments.formatter import Formatter
+
+_T = TypeVar("_T", str, bytes)
 
 def lex(code, lexer): ...
-def format(tokens, formatter, outfile: Any | None = ...): ...
-def highlight(code, lexer, formatter, outfile: Any | None = ...): ...
+@overload
+def format(tokens, formatter: Formatter[_T], outfile: SupportsWrite[_T]) -> None: ...
+@overload
+def format(tokens, formatter: Formatter[_T], outfile: None = ...) -> _T: ...
+@overload
+def highlight(code, lexer, formatter: Formatter[_T], outfile: SupportsWrite[_T]) -> None: ...
+@overload
+def highlight(code, lexer, formatter: Formatter[_T], outfile: None = ...) -> _T: ...

stubs/Pygments/pygments/formatter.pyi

@@ -1,6 +1,8 @@
-from typing import Any
+from typing import Any, Generic, TypeVar, overload
 
-class Formatter:
+_T = TypeVar("_T", str, bytes)
+
+class Formatter(Generic[_T]):
     name: Any
     aliases: Any
     filenames: Any
@@ -10,6 +12,11 @@ class Formatter:
     title: Any
     encoding: Any
     options: Any
-    def __init__(self, **options) -> None: ...
+    @overload
+    def __init__(self: Formatter[str], *, encoding: None = ..., outencoding: None = ..., **options) -> None: ...
+    @overload
+    def __init__(self: Formatter[bytes], *, encoding: str, outencoding: None = ..., **options) -> None: ...
+    @overload
+    def __init__(self: Formatter[bytes], *, encoding: None = ..., outencoding: str, **options) -> None: ...
     def get_style_defs(self, arg: str = ...): ...
     def format(self, tokensource, outfile): ...

stubs/Pygments/pygments/formatters/__init__.pyi

@@ -1,4 +1,4 @@
-from typing import Generator
+from typing import Any, Generator
 
 from ..formatter import Formatter
 from .bbcode import BBCodeFormatter as BBCodeFormatter
@@ -18,7 +18,7 @@ from .svg import SvgFormatter as SvgFormatter
 from .terminal import TerminalFormatter as TerminalFormatter
 from .terminal256 import Terminal256Formatter as Terminal256Formatter, TerminalTrueColorFormatter as TerminalTrueColorFormatter
 
-def get_all_formatters() -> Generator[type[Formatter], None, None]: ...
+def get_all_formatters() -> Generator[type[Formatter[Any]], None, None]: ...
 def get_formatter_by_name(_alias, **options): ...
 def load_formatter_from_file(filename, formattername: str = ..., **options): ...
 def get_formatter_for_filename(fn, **options): ...

stubs/Pygments/pygments/formatters/bbcode.pyi

@@ -1,11 +1,12 @@
-from typing import Any
+from typing import Any, TypeVar
 
 from pygments.formatter import Formatter
 
-class BBCodeFormatter(Formatter):
+_T = TypeVar("_T", str, bytes)
+
+class BBCodeFormatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
     styles: Any
-    def __init__(self, **options) -> None: ...
     def format_unencoded(self, tokensource, outfile) -> None: ...

stubs/Pygments/pygments/formatters/html.pyi

@@ -1,8 +1,10 @@
-from typing import Any
+from typing import Any, TypeVar
 
 from pygments.formatter import Formatter
 
-class HtmlFormatter(Formatter):
+_T = TypeVar("_T", str, bytes)
+
+class HtmlFormatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
@@ -30,7 +32,6 @@ class HtmlFormatter(Formatter):
     linespans: Any
     anchorlinenos: Any
     hl_lines: Any
-    def __init__(self, **options) -> None: ...
     def get_style_defs(self, arg: Any | None = ...): ...
     def get_token_style_defs(self, arg: Any | None = ...): ...
     def get_background_style_defs(self, arg: Any | None = ...): ...

stubs/Pygments/pygments/formatters/img.pyi

@@ -1,7 +1,9 @@
-from typing import Any
+from typing import Any, TypeVar
 
 from pygments.formatter import Formatter
 
+_T = TypeVar("_T", str, bytes)
+
 class PilNotAvailable(ImportError): ...
 class FontNotFound(Exception): ...
 
@@ -15,7 +17,7 @@ class FontManager:
     def get_text_size(self, text): ...
     def get_font(self, bold, oblique): ...
 
-class ImageFormatter(Formatter):
+class ImageFormatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
@@ -42,23 +44,22 @@ class ImageFormatter(Formatter):
     hl_lines: Any
     hl_color: Any
     drawables: Any
-    def __init__(self, **options) -> None: ...
     def get_style_defs(self, arg: str = ...) -> None: ...
     def format(self, tokensource, outfile) -> None: ...
 
-class GifImageFormatter(ImageFormatter):
+class GifImageFormatter(ImageFormatter[_T]):
     name: str
     aliases: Any
     filenames: Any
     default_image_format: str
 
-class JpgImageFormatter(ImageFormatter):
+class JpgImageFormatter(ImageFormatter[_T]):
     name: str
     aliases: Any
     filenames: Any
     default_image_format: str
 
-class BmpImageFormatter(ImageFormatter):
+class BmpImageFormatter(ImageFormatter[_T]):
     name: str
     aliases: Any
     filenames: Any

stubs/Pygments/pygments/formatters/irc.pyi

@@ -1,13 +1,14 @@
-from typing import Any
+from typing import Any, TypeVar
 
 from pygments.formatter import Formatter
 
-class IRCFormatter(Formatter):
+_T = TypeVar("_T", str, bytes)
+
+class IRCFormatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
     darkbg: Any
     colorscheme: Any
     linenos: Any
-    def __init__(self, **options) -> None: ...
     def format_unencoded(self, tokensource, outfile) -> None: ...

stubs/Pygments/pygments/formatters/latex.pyi

@@ -1,9 +1,11 @@
-from typing import Any
+from typing import Any, TypeVar
 
 from pygments.formatter import Formatter
 from pygments.lexer import Lexer
 
-class LatexFormatter(Formatter):
+_T = TypeVar("_T", str, bytes)
+
+class LatexFormatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
@@ -21,7 +23,6 @@ class LatexFormatter(Formatter):
     left: Any
     right: Any
     envname: Any
-    def __init__(self, **options) -> None: ...
     def get_style_defs(self, arg: str = ...): ...
     def format_unencoded(self, tokensource, outfile) -> None: ...
 

stubs/Pygments/pygments/formatters/other.pyi

@@ -1,26 +1,26 @@
-from typing import Any
+from typing import Any, TypeVar
 
 from pygments.formatter import Formatter
 
-class NullFormatter(Formatter):
+_T = TypeVar("_T", str, bytes)
+
+class NullFormatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
     def format(self, tokensource, outfile) -> None: ...
 
-class RawTokenFormatter(Formatter):
+class RawTokenFormatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
     unicodeoutput: bool
     encoding: str
     compress: Any
     error_color: Any
-    def __init__(self, **options) -> None: ...
     def format(self, tokensource, outfile) -> None: ...
 
-class TestcaseFormatter(Formatter):
+class TestcaseFormatter(Formatter[_T]):
     name: str
     aliases: Any
-    def __init__(self, **options) -> None: ...
     def format(self, tokensource, outfile) -> None: ...

stubs/Pygments/pygments/formatters/pangomarkup.pyi

@@ -1,11 +1,12 @@
-from typing import Any
+from typing import Any, TypeVar
 
 from pygments.formatter import Formatter
 
-class PangoMarkupFormatter(Formatter):
+_T = TypeVar("_T", str, bytes)
+
+class PangoMarkupFormatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
     styles: Any
-    def __init__(self, **options) -> None: ...
     def format_unencoded(self, tokensource, outfile) -> None: ...

stubs/Pygments/pygments/formatters/rtf.pyi

@@ -1,12 +1,13 @@
-from typing import Any
+from typing import Any, TypeVar
 
 from pygments.formatter import Formatter
 
-class RtfFormatter(Formatter):
+_T = TypeVar("_T", str, bytes)
+
+class RtfFormatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
     fontface: Any
     fontsize: Any
-    def __init__(self, **options) -> None: ...
     def format_unencoded(self, tokensource, outfile) -> None: ...

stubs/Pygments/pygments/formatters/svg.pyi

@@ -1,8 +1,10 @@
-from typing import Any
+from typing import Any, TypeVar
 
 from pygments.formatter import Formatter
 
-class SvgFormatter(Formatter):
+_T = TypeVar("_T", str, bytes)
+
+class SvgFormatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
@@ -17,5 +19,4 @@ class SvgFormatter(Formatter):
     linenostart: Any
     linenostep: Any
     linenowidth: Any
-    def __init__(self, **options) -> None: ...
     def format_unencoded(self, tokensource, outfile) -> None: ...

stubs/Pygments/pygments/formatters/terminal.pyi

@@ -1,14 +1,15 @@
-from typing import Any
+from typing import Any, TypeVar
 
 from pygments.formatter import Formatter
 
-class TerminalFormatter(Formatter):
+_T = TypeVar("_T", str, bytes)
+
+class TerminalFormatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
     darkbg: Any
     colorscheme: Any
     linenos: Any
-    def __init__(self, **options) -> None: ...
     def format(self, tokensource, outfile): ...
     def format_unencoded(self, tokensource, outfile) -> None: ...

stubs/Pygments/pygments/formatters/terminal256.pyi

@@ -1,7 +1,9 @@
-from typing import Any
+from typing import Any, TypeVar
 
 from pygments.formatter import Formatter
 
+_T = TypeVar("_T", str, bytes)
+
 class EscapeSequence:
     fg: Any
     bg: Any
@@ -16,7 +18,7 @@ class EscapeSequence:
     def true_color_string(self): ...
     def reset_string(self): ...
 
-class Terminal256Formatter(Formatter):
+class Terminal256Formatter(Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any
@@ -27,11 +29,10 @@ class Terminal256Formatter(Formatter):
     useunderline: Any
     useitalic: Any
     linenos: Any
-    def __init__(self, **options) -> None: ...
     def format(self, tokensource, outfile): ...
     def format_unencoded(self, tokensource, outfile) -> None: ...
 
-class TerminalTrueColorFormatter(Terminal256Formatter):
+class TerminalTrueColorFormatter(Terminal256Formatter[_T]):
     name: str
     aliases: Any
     filenames: Any

stubs/Pygments/pygments/plugin.pyi

@@ -1,4 +1,4 @@
-from typing import Generator, Iterable
+from typing import Any, Generator, Iterable
 
 from pkg_resources import EntryPoint
 from pygments.filter import Filter
@@ -13,6 +13,6 @@ FILTER_ENTRY_POINT: str
 
 def iter_entry_points(group_name: str) -> Iterable[EntryPoint]: ...
 def find_plugin_lexers() -> Generator[type[Lexer], None, None]: ...
-def find_plugin_formatters() -> Generator[tuple[str, type[Formatter]], None, None]: ...
+def find_plugin_formatters() -> Generator[tuple[str, type[Formatter[Any]]], None, None]: ...
 def find_plugin_styles() -> Generator[tuple[str, type[Style]], None, None]: ...
 def find_plugin_filters() -> Generator[tuple[str, type[Filter]], None, None]: ...