Add Type Hints

pyproject.toml

@@ -19,6 +19,7 @@ classifiers = [
 dependencies = [
     "docopt>=0.6.2",
 ]
+requires-python = ">=3.8"
 
 [project.scripts]
 xortool = 'xortool.tool_main:main'
@@ -29,6 +30,7 @@ homepage = "http://github.com/hellman/xortool"
 
 [tool.poetry.group.dev.dependencies]
 importlib_metadata = "^4.8"
+types-docopt = "^0.6"
 
 [build-system]
 requires = ["poetry-core>=2.0.0,<3.0.0"]

xortool/args.py

@@ -1,13 +1,37 @@
+from __future__ import annotations
+
+from typing import Literal, overload, TypedDict
 from docopt import docopt
 
 from xortool.charset import get_charset
 
+class ParameterDict(TypedDict):
+
+    brute_chars: bool
+    brute_printable: bool
+    filename: str
+    filter_output: bool
+    frequency_spread: int
+    input_is_hex: bool
+    known_key_length: int | None
+    max_key_length: int | None
+    most_frequent_char: int | None
+    text_charset: str | bytes
+    known_plain: bytes | Literal[False]
+    threshold: int | None
+
+
 
 class ArgError(Exception):
     pass
 
+@overload
+def parse_char(ch: None) -> None: ...
+
+@overload
+def parse_char(ch: str) -> int: ...
 
-def parse_char(ch):
+def parse_char(ch: str | None) -> int | None:
     """
     'A' or '\x41' or '0x41' or '41'
     '\x00' or '0x00' or '00'
@@ -24,14 +48,19 @@ def parse_char(ch):
         raise ValueError("Char can be only a char letter or hex")
     return int(ch, 16)
 
+@overload
+def parse_int(i: None) -> None: ...
+
+@overload
+def parse_int(i: str) -> int: ...
 
-def parse_int(i):
+def parse_int(i: str | None) -> int | None:
     if i is None:
         return None
     return int(i)
 
 
-def parse_parameters(doc, version):
+def parse_parameters(doc: str, version: str) -> ParameterDict:
     p = docopt(doc, version=version)
     p = {k.lstrip("-"): v for k, v in p.items()}
     try:

xortool/charset.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import string
 
 
@@ -20,7 +22,7 @@ class CharsetError(Exception):
 }
 
 
-def get_charset(charset):
+def get_charset(charset: str | None) -> str | bytes:
     charset = charset or "printable"
     if charset in PREDEFINED_CHARSETS:
         return PREDEFINED_CHARSETS[charset].encode("ascii")

xortool/libcolors.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import os
+from typing import cast
 
 
 BASH_ATTRIBUTES = {"regular": "0",
@@ -13,7 +16,7 @@
                  "blue": "44", "purple": "45", "cyan": "46", "white": "47"}
 
 
-def _main():
+def _main() -> None:
     header = color("white", "black", "dark")
     print()
 
@@ -36,7 +39,7 @@ def _main():
     print()
 
 
-def color(color=None, bgcolor=None, attrs=None):
+def color(color: str | None = None, bgcolor: str | None = None, attrs: str | None = None) -> str:
     if not is_bash():
         return ""
 
@@ -61,13 +64,15 @@ def color(color=None, bgcolor=None, attrs=None):
     return ret + "m"
 
 
-def is_bash():
+def is_bash() -> bool:
     return os.environ.get("SHELL", "unknown").endswith("bash")
 
 
-def _keys_sorted_by_values(adict):
+def _keys_sorted_by_values(adict: dict[str, str]) -> list[str]:
     """Return list of the keys of @adict sorted by values."""
-    return sorted(adict, key=adict.get)
+    # Casting is fine here, we are sorting, adict.get will always
+    # get a value, and never return "None"
+    return sorted(adict, key=lambda v: cast(str, adict.get(v)))
 
 
 if __name__ == "__main__":

xortool/routine.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import sys
 import string
@@ -7,19 +9,19 @@ class MkdirError(Exception):
     pass
 
 
-def load_file(filename):
+def load_file(filename: str | int) -> bytes:
     if filename == "-":
         filename = sys.stdin.fileno()
     with open(filename, "rb") as fd:
         return fd.read()
 
 
-def save_file(filename, data):
+def save_file(filename: str, data: bytes) -> None:
     with open(filename, "wb") as fd:
         fd.write(data)
 
 
-def mkdir(dirname):
+def mkdir(dirname: str) -> None:
     if os.path.exists(dirname):
         return
     try:
@@ -28,7 +30,7 @@ def mkdir(dirname):
         raise MkdirError(str(err))
 
 
-def rmdir(dirname):
+def rmdir(dirname: str) -> None:
     if dirname[-1] == os.sep:
         dirname = dirname[:-1]
     if os.path.islink(dirname):
@@ -43,30 +45,30 @@ def rmdir(dirname):
             os.unlink(path)
     os.rmdir(dirname)
 
-def decode_from_hex(text):
-    text = text.decode(encoding='ascii', errors='ignore')
-    only_hex_digits = "".join(c for c in text if c in string.hexdigits)
+def decode_from_hex(text: bytes) -> bytes:
+    text_bytes = text.decode(encoding='ascii', errors='ignore')
+    only_hex_digits = "".join(c for c in text_bytes if c in string.hexdigits)
     return bytes.fromhex(only_hex_digits)
 
 
-def dexor(text, key):
+def dexor(text: bytes, key: bytes) -> bytes:
     mod = len(key)
     return bytes(key[index % mod] ^ char for index, char in enumerate(text))
 
 
-def die(exitMessage, exitCode=1):
+def die(exitMessage: str, exitCode: int = 1) -> None:
     print(exitMessage)
     sys.exit(exitCode)
 
 
-def is_linux():
+def is_linux() -> bool:
     return sys.platform.startswith("linux")
 
 
-def alphanum(s):
+def alphanum(s: str) -> str:
     lst = list(s)
     for index, char in enumerate(lst):
         if char in string.ascii_letters + string.digits:
             continue
-        lst[index] = char.hex()
+        lst[index] = char.encode("ascii").hex()
     return "".join(lst)