Add type hinting (#143)

Co-authored-by: Sorin Sbarnea <[email protected]>

Author: ziegenberg

.pre-commit-config.yaml

@@ -37,3 +37,19 @@ repos:
         additional_dependencies:
           - ansible-base
           - testinfra
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v0.910-1
+    hooks:
+    - id: mypy
+      # empty args needed in order to match mypy cli behavior
+      args: ["--strict"]
+      additional_dependencies:
+      - pytest
+      - types-setuptools
+      - types-pkg_resources
+      - types-mock
+      exclude: >
+        (?x)^(
+          docs/.*|
+          setup.py
+        )$

ansi2html/converter.py

@@ -20,17 +20,17 @@
 #  along with this program.  If not, see
 #  <http://www.gnu.org/licenses/>.
 
-import io
 import optparse
 import re
 import sys
+from typing import Dict, Iterator, List, Optional, Set, Tuple, Union
 
 import pkg_resources
 
 try:
     from collections import OrderedDict
 except ImportError:
-    from ordereddict import OrderedDict
+    from ordereddict import OrderedDict  # type: ignore
 
 from ansi2html.style import SCHEME, get_styles
 
@@ -115,22 +115,22 @@
 
 
 class _State:
-    def __init__(self):
+    def __init__(self) -> None:
         self.inside_span = False
         self.reset()
 
-    def reset(self):
-        self.intensity = ANSI_INTENSITY_NORMAL
-        self.style = ANSI_STYLE_NORMAL
-        self.blink = ANSI_BLINK_OFF
-        self.underline = ANSI_UNDERLINE_OFF
-        self.crossedout = ANSI_CROSSED_OUT_OFF
-        self.visibility = ANSI_VISIBILITY_ON
-        self.foreground = (ANSI_FOREGROUND_DEFAULT, None)
-        self.background = (ANSI_BACKGROUND_DEFAULT, None)
-        self.negative = ANSI_NEGATIVE_OFF
-
-    def adjust(self, ansi_code, parameter=None):
+    def reset(self) -> None:
+        self.intensity: int = ANSI_INTENSITY_NORMAL
+        self.style: int = ANSI_STYLE_NORMAL
+        self.blink: int = ANSI_BLINK_OFF
+        self.underline: int = ANSI_UNDERLINE_OFF
+        self.crossedout: int = ANSI_CROSSED_OUT_OFF
+        self.visibility: int = ANSI_VISIBILITY_ON
+        self.foreground: Tuple[int, Optional[int]] = (ANSI_FOREGROUND_DEFAULT, None)
+        self.background: Tuple[int, Optional[int]] = (ANSI_BACKGROUND_DEFAULT, None)
+        self.negative: int = ANSI_NEGATIVE_OFF
+
+    def adjust(self, ansi_code: int, parameter: Optional[int] = None) -> None:
         if ansi_code in (
             ANSI_INTENSITY_INCREASED,
             ANSI_INTENSITY_REDUCED,
@@ -174,17 +174,21 @@ def adjust(self, ansi_code, parameter=None):
         elif ansi_code in (ANSI_NEGATIVE_ON, ANSI_NEGATIVE_OFF):
             self.negative = ansi_code
 
-    def to_css_classes(self):
-        css_classes = []
+    def to_css_classes(self) -> List[str]:
+        css_classes: List[str] = []
 
-        def append_unless_default(output, value, default):
+        def append_unless_default(output: List[str], value: int, default: int) -> None:
             if value != default:
                 css_class = "ansi%d" % value
                 output.append(css_class)
 
         def append_color_unless_default(
-            output, color, default, negative, neg_css_class
-        ):
+            output: List[str],
+            color: Tuple[int, Optional[int]],
+            default: int,
+            negative: bool,
+            neg_css_class: str,
+        ) -> None:
             value, parameter = color
             if value != default:
                 prefix = "inv" if negative else "ansi"
@@ -222,17 +226,17 @@ def append_color_unless_default(
 
 
 class OSC_Link:
-    def __init__(self, url, text):
+    def __init__(self, url: str, text: str) -> None:
         self.url = url
         self.text = text
 
 
-def map_vt100_box_code(char):
+def map_vt100_box_code(char: str) -> str:
     char_hex = hex(ord(char))
     return VT100_BOX_CODES[char_hex] if char_hex in VT100_BOX_CODES else char
 
 
-def _needs_extra_newline(text):
+def _needs_extra_newline(text: str) -> bool:
     if not text or text.endswith("\n"):
         return False
     return True
@@ -254,18 +258,18 @@ class Ansi2HTMLConverter:
 
     def __init__(
         self,
-        latex=False,
-        inline=False,
-        dark_bg=True,
-        line_wrap=True,
-        font_size="normal",
-        linkify=False,
-        escaped=True,
-        markup_lines=False,
-        output_encoding="utf-8",
-        scheme="ansi2html",
-        title="",
-    ):
+        latex: bool = False,
+        inline: bool = False,
+        dark_bg: bool = True,
+        line_wrap: bool = True,
+        font_size: str = "normal",
+        linkify: bool = False,
+        escaped: bool = True,
+        markup_lines: bool = False,
+        output_encoding: str = "utf-8",
+        scheme: str = "ansi2html",
+        title: str = "",
+    ) -> None:
 
         self.latex = latex
         self.inline = inline
@@ -278,7 +282,7 @@ def __init__(
         self.output_encoding = output_encoding
         self.scheme = scheme
         self.title = title
-        self._attrs = None
+        self._attrs: Optional[Dict[str, Union[bool, str, Set[str]]]] = None
         self.hyperref = False
 
         if inline:
@@ -298,27 +302,27 @@ def __init__(
         )
         self.osc_link_re = re.compile("\033\\]8;;(.*?)\007(.*?)\033\\]8;;\007")
 
-    def do_linkify(self, line):
+    def do_linkify(self, line: str) -> str:
         if not isinstance(line, str):
             return line  # If line is an object, e.g. OSC_Link, it
             # will be expanded to a string later
         if self.latex:
             return self.url_matcher.sub(r"\\url{\1}", line)
         return self.url_matcher.sub(r'<a href="\1">\1</a>', line)
 
-    def handle_osc_links(self, part):
+    def handle_osc_links(self, part: OSC_Link) -> str:
         if self.latex:
             self.hyperref = True
             return """\\href{%s}{%s}""" % (part.url, part.text)
         return """<a href="%s">%s</a>""" % (part.url, part.text)
 
-    def apply_regex(self, ansi):
-        styles_used = set()
-        parts = self._apply_regex(ansi, styles_used)
-        parts = self._collapse_cursor(parts)
-        parts = list(parts)
+    def apply_regex(self, ansi: str) -> Tuple[str, Set[str]]:
+        styles_used: Set[str] = set()
+        all_parts = self._apply_regex(ansi, styles_used)
+        no_cursor_parts = self._collapse_cursor(all_parts)
+        no_cursor_parts = list(no_cursor_parts)
 
-        def _check_links(parts):
+        def _check_links(parts: List[Union[str, OSC_Link]]) -> Iterator[str]:
             for part in parts:
                 if isinstance(part, str):
                     if self.linkify:
@@ -330,7 +334,7 @@ def _check_links(parts):
                 else:
                     yield part
 
-        parts = list(_check_links(parts))
+        parts = list(_check_links(no_cursor_parts))
         combined = "".join(parts)
         if self.markup_lines and not self.latex:
             combined = "\n".join(
@@ -341,7 +345,9 @@ def _check_links(parts):
             )
         return combined, styles_used
 
-    def _apply_regex(self, ansi, styles_used):
+    def _apply_regex(
+        self, ansi: str, styles_used: Set[str]
+    ) -> Iterator[Union[str, OSC_Link, CursorMoveUp]]:
         if self.escaped:
             if (
                 self.latex
@@ -358,7 +364,7 @@ def _apply_regex(self, ansi, styles_used):
             for pattern, special in specials.items():
                 ansi = ansi.replace(pattern, special)
 
-        def _vt100_box_drawing():
+        def _vt100_box_drawing() -> Iterator[str]:
             last_end = 0  # the index of the last end of a code we've seen
             box_drawing_mode = False
             for match in self.vt100_box_codes_prog.finditer(ansi):
@@ -374,7 +380,7 @@ def _vt100_box_drawing():
 
         ansi = "".join(_vt100_box_drawing())
 
-        def _osc_link(ansi):
+        def _osc_link(ansi: str) -> Iterator[Union[str, OSC_Link]]:
             last_end = 0
             for match in self.osc_link_re.finditer(ansi):
                 trailer = ansi[last_end : match.start()]
@@ -397,20 +403,23 @@ def _osc_link(ansi):
             else:
                 yield "</span>"
 
-    def _handle_ansi_code(self, ansi, styles_used, state):
+    def _handle_ansi_code(
+        self, ansi: str, styles_used: Set[str], state: _State
+    ) -> Iterator[Union[str, CursorMoveUp]]:
         last_end = 0  # the index of the last end of a code we've seen
         for match in self.ansi_codes_prog.finditer(ansi):
             yield ansi[last_end : match.start()]
             last_end = match.end()
 
+            params: Union[str, List[int]]
             params, command = match.groups()
 
             if command not in "mMA":
                 continue
 
             # Special cursor-moving code.  The only supported one.
             if command == "A":
-                yield CursorMoveUp
+                yield CursorMoveUp  # type: ignore
                 continue
 
             try:
@@ -452,7 +461,7 @@ def _handle_ansi_code(self, ansi, styles_used, state):
 
                 if v in (ANSI_FOREGROUND_256, ANSI_BACKGROUND_256):
                     try:
-                        parameter = params[i + 2]
+                        parameter: Optional[int] = params[i + 2]
                     except IndexError:
                         continue
                     skip_after_index = i + 2
@@ -495,22 +504,24 @@ def _handle_ansi_code(self, ansi, styles_used, state):
             state.inside_span = True
         yield ansi[last_end:]
 
-    def _collapse_cursor(self, parts):
+    def _collapse_cursor(
+        self, parts: Iterator[Union[str, OSC_Link, CursorMoveUp]]
+    ) -> List[Union[str, OSC_Link]]:
         """Act on any CursorMoveUp commands by deleting preceding tokens"""
 
-        final_parts = []
+        final_parts: List[Union[str, OSC_Link]] = []
         for part in parts:
 
             # Throw out empty string tokens ("")
             if not part:
                 continue
 
             # Go back, deleting every token in the last 'line'
-            if part == CursorMoveUp:
+            if isinstance(part, CursorMoveUp):
                 if final_parts:
                     final_parts.pop()
 
-                while final_parts and "\n" not in final_parts[-1]:
+                while final_parts and "\n" not in final_parts[-1]:  # type: ignore
                     final_parts.pop()
 
                 continue
@@ -520,7 +531,9 @@ def _collapse_cursor(self, parts):
 
         return final_parts
 
-    def prepare(self, ansi="", ensure_trailing_newline=False):
+    def prepare(
+        self, ansi: str = "", ensure_trailing_newline: bool = False
+    ) -> Dict[str, Union[bool, str, Set[str]]]:
         """Load the contents of 'ansi' into this object"""
 
         body, styles = self.apply_regex(ansi)
@@ -538,29 +551,31 @@ def prepare(self, ansi="", ensure_trailing_newline=False):
 
         return self._attrs
 
-    def attrs(self):
+    def attrs(self) -> Dict[str, Union[bool, str, Set[str]]]:
         """Prepare attributes for the template"""
         if not self._attrs:
             raise Exception("Method .prepare not yet called.")
         return self._attrs
 
-    def convert(self, ansi, full=True, ensure_trailing_newline=False):
+    def convert(
+        self, ansi: str, full: bool = True, ensure_trailing_newline: bool = False
+    ) -> str:
         r"""
         :param ansi: ANSI sequence to convert.
         :param full: Whether to include the full HTML document or only the body.
         :param ensure_trailing_newline: Ensures that ``\n`` character is present at the end of the output.
         """
         attrs = self.prepare(ansi, ensure_trailing_newline=ensure_trailing_newline)
         if not full:
-            return attrs["body"]
+            return attrs["body"]  # type: ignore
         if self.latex:
             _template = _latex_template
         else:
             _template = _html_template
         all_styles = get_styles(self.dark_bg, self.line_wrap, self.scheme)
         backgrounds = all_styles[:6]
         used_styles = filter(
-            lambda e: e.klass.lstrip(".") in attrs["styles"], all_styles
+            lambda e: e.klass.lstrip(".") in attrs["styles"], all_styles  # type: ignore
         )
 
         return _template % {
@@ -572,15 +587,15 @@ def convert(self, ansi, full=True, ensure_trailing_newline=False):
             "hyperref": "\\usepackage{hyperref}" if self.hyperref else "",
         }
 
-    def produce_headers(self):
+    def produce_headers(self) -> str:
         return '<style type="text/css">\n%(style)s\n</style>\n' % {
             "style": "\n".join(
                 map(str, get_styles(self.dark_bg, self.line_wrap, self.scheme))
             )
         }
 
 
-def main():
+def main() -> None:
     """
     $ ls --color=always | ansi2html > directories.html
     $ sudo tail /var/log/messages | ccze -A | ansi2html > logs.html
@@ -718,16 +733,7 @@ def main():
         title=opts.output_title,
     )
 
-    try:
-        sys.stdin = io.TextIOWrapper(sys.stdin.detach(), opts.input_encoding, "replace")
-    except io.UnsupportedOperation:
-        # This only fails in the test suite...
-        pass
-
-    def _read(input_bytes):
-        return input_bytes
-
-    def _print(output_unicode, end="\n"):
+    def _print(output_unicode: str, end: str = "\n") -> None:
         if hasattr(sys.stdout, "buffer"):
             output_bytes = (output_unicode + end).encode(opts.output_encoding)
             sys.stdout.buffer.write(output_bytes)