Add type hints for all examples

niklasf · niklasf · commit 6e09f0ff5316 · 2020-08-10T12:19:50.000+02:00
diff --git a/.travis.yml b/.travis.yml
@@ -86,6 +86,11 @@ script:
  - if [[ $LATEST -eq 1 ]]; then python -m mypy --strict chess/syzygy.py; fi
  - if [[ $LATEST -eq 1 ]]; then python -m mypy --strict chess/variant.py; fi
  - # Typing of examples
+ - if [[ $LATEST -eq 1 ]]; then python -m mypy --strict examples/bratko_kopec/bratko_kopec.py || true; fi
+ - if [[ $LATEST -eq 1 ]]; then python -m mypy --strict examples/chess960_pos_list.py; fi
+ - if [[ $LATEST -eq 1 ]]; then python -m mypy examples/perft/perft.py; fi
+ - if [[ $LATEST -eq 1 ]]; then python -m mypy examples/polyglot_tree.py; fi
+ - if [[ $LATEST -eq 1 ]]; then python -m mypy --strict examples/push_san.py; fi
  - if [[ $LATEST -eq 1 ]]; then python -m mypy --strict examples/xray_attacks.py; fi
  - # Perft tests
  - if [[ $LATEST -eq 1 ]]; then python examples/perft/perft.py -t 1 examples/perft/random.perft --max-nodes 10000; fi
diff --git a/examples/bratko_kopec/bratko_kopec.py b/examples/bratko_kopec/bratko_kopec.py
@@ -9,12 +9,14 @@
 import logging
 import sys
 
+from typing import Type
+
 import chess
 import chess.engine
 import chess.variant
 
 
-async def test_epd(engine, epd, VariantBoard, movetime):
+async def test_epd(engine: chess.engine.EngineProtocol, epd: str, VariantBoard: Type[chess.Board], movetime: float) -> float:
     board, epd_info = VariantBoard.from_epd(epd)
     epd_string = epd_info.get("id", board.fen())
     if "am" in epd_info:
@@ -25,7 +27,10 @@ async def test_epd(engine, epd, VariantBoard, movetime):
     limit = chess.engine.Limit(time=movetime)
     result = await engine.play(board, limit, game=object())
 
-    if "am" in epd_info and result.move in epd_info["am"]:
+    if not result.move:
+        print(f"{epd_string}: -- | +0")
+        return 0.0
+    elif "am" in epd_info and result.move in epd_info["am"]:
         print(f"{epd_string}: {board.san(result.move)} | +0")
         return 0.0
     elif "bm" in epd_info and result.move not in epd_info["bm"]:
@@ -36,7 +41,7 @@ async def test_epd(engine, epd, VariantBoard, movetime):
         return 1.0
 
 
-async def test_epd_with_fractional_scores(engine, epd, VariantBoard, movetime):
+async def test_epd_with_fractional_scores(engine: chess.engine.EngineProtocol, epd: str, VariantBoard: Type[chess.Board], movetime: float) -> float:
     board, epd_info = VariantBoard.from_epd(epd)
     epd_string = epd_info.get("id", board.fen())
     if "am" in epd_info:
@@ -71,7 +76,7 @@ async def test_epd_with_fractional_scores(engine, epd, VariantBoard, movetime):
     return score
 
 
-async def main():
+async def main() -> None:
     # Parse command line arguments.
     parser = argparse.ArgumentParser(description=__doc__)
 
diff --git a/examples/chess960_pos_list.py b/examples/chess960_pos_list.py
@@ -7,7 +7,7 @@
 import chess
 
 
-def main(bench_only=False):
+def main(bench_only: bool = False) -> None:
     board = chess.Board.empty(chess960=True)
 
     for sharnagl in range(0, 960):
diff --git a/examples/perft/perft.py b/examples/perft/perft.py
@@ -4,16 +4,19 @@
 Run perft test to check correctness and speed of the legal move generator.
 """
 
-import chess
-import chess.variant
 import multiprocessing
 import functools
 import time
 import argparse
 import sys
 
+from typing import Callable, Iterator, Optional, TextIO, Type
+
+import chess
+import chess.variant
+
 
-def perft(depth, board):
+def perft(depth: int, board: chess.Board) -> int:
     if depth == 1:
         return board.legal_moves.count()
     elif depth > 1:
@@ -29,11 +32,11 @@ def perft(depth, board):
         return 1
 
 
-def parallel_perft(pool, depth, board):
+def parallel_perft(pool, depth: int, board: chess.Board) -> int:
     if depth == 1:
         return board.legal_moves.count()
     elif depth > 1:
-        def successors(board):
+        def successors(board: chess.Board) -> Iterator[chess.Board]:
             for move in board.legal_moves:
                 board_after = board.copy(stack=False)
                 board_after.push(move)
@@ -44,14 +47,14 @@ def successors(board):
         return 1
 
 
-def sdiv(a, b):
+def sdiv(a: float, b: float) -> float:
     try:
         return a / b
     except ZeroDivisionError:
         return float("Inf")
 
 
-def main(perft_file, VariantBoard, perft_f, max_depth, max_nodes):
+def main(perft_file: TextIO, VariantBoard: Type[chess.Board], perft_f: Callable[[int, chess.Board], int], max_depth: Optional[int], max_nodes: Optional[int]) -> None:
     current_id = None
     board = VariantBoard(chess960=True)
     column = 0
diff --git a/examples/polyglot_tree.py b/examples/polyglot_tree.py
@@ -2,12 +2,15 @@
 
 """Print a Polyglot opening book in tree form."""
 
+import argparse
+
+from typing import Set
+
 import chess
 import chess.polyglot
-import argparse
 
 
-def print_tree(args, visited, level=0):
+def print_tree(args: argparse.Namespace, visited: Set[int], level: int = 0) -> None:
     if level >= args.depth:
         return
 
diff --git a/examples/push_san.py b/examples/push_san.py
@@ -6,7 +6,7 @@
 import timeit
 
 
-def play_immortal_game():
+def play_immortal_game() -> None:
     board = chess.Board()
 
     # 1. e4 e5
