Added new optional filters to main script

README.md

@@ -83,6 +83,11 @@ To execute the generator execute this command. By default it will look for the `
 - `--includeBlunder=False` If False then generated puzzles won't include initial blunder move, default is `True`
 - `--stockfish=./stockfish-x86_64-bmi2` Path to Stockfish binary. 
   Optional. If omitted, the program will try to locate Stockfish in current directory or download it from the net
+- `--color=WHITE` Generate puzzles only for preferred color. `WHITE`/`W`, `BLACK`/`B`, or default `NONE` 
+- `--opening= '1. e4 e5 2. f4` Generate puzzles stemming from opening. Default is `NONE`
+- `--forcedMate=TRUE` Only generate puzzles with a forced checkmate. Default is `False`
+- `--minTurn=10` Generate puzzles from turns >= minTurn. Default is `0` 
+- `--maxTurn=50` Generate puzzles from turns <= maxTurn. Default is `999`
 
 Example:
 `python3 main.py --quiet --depth=12 --games=ruy_lopez.pgn --strict=True --threads=2 --memory=1024`

main.py

@@ -4,13 +4,16 @@
 
 import argparse
 import logging
+import io
+import sys
 
 import chess.engine
 import chess.pgn
 
 from modules.api.api import post_puzzle
 from modules.bcolors.bcolors import bcolors
 from modules.investigate.investigate import investigate
+from modules.investigate.investigate import filter_game
 from modules.puzzle.puzzle import puzzle
 from modules.utils.helpers import str2bool, get_stockfish_command, configure_logging, prepare_terminal
 
@@ -36,6 +39,21 @@ def prepare_settings():
                         help="If False then generated puzzles won't include initial blunder move")
     parser.add_argument("--stockfish", metavar="STOCKFISH", default=None, help="Path to Stockfish binary")
 
+    parser.add_argument("--color", metavar="DESIRED_COLOR", type=str, default="EITHER",
+                        help="desired color in puzzle- BLACK, WHITE, or EITHER(default)")
+
+    parser.add_argument("--forcedMate", metavar="MATES_ONLY", type=bool, default=False,
+                        help="only return forced checkmates")
+
+    parser.add_argument("--minTurn", metavar="MIN_TURN", type=int, default=0,
+                        help="only return tactics starting after inputted turn")
+
+    parser.add_argument("--maxTurn", metavar="MAX_TURN", type=int, default=999,
+                        help="only return tactics starting before inputted turn")
+
+    parser.add_argument("--opening", metavar="ROOT_PGN", type=str, default="NONE",
+                        help="to get tactics from a specific opening, input opening moves as pgn")
+
     return parser.parse_args()
 
 
@@ -53,10 +71,35 @@ def prepare_settings():
 all_games = open(settings.games, "r")
 tactics_file = open("tactics.pgn", "w")
 game_id = 0
+
+opening_str = settings.opening
+checkmate_only = settings.forcedMate
+min_turn = settings.minTurn
+max_turn = settings.maxTurn
+color_str = settings.color.upper()
+preferred_color = 0 if (color_str == "BLACK" or color_str == 'B') else 1 if (
+            color_str == "WHITE" or color_str == "W") else 2
+opening_pgn = io.StringIO(settings.opening)
+opening = chess.pgn.read_game(opening_pgn)
+
+if len(opening.errors) > 0:
+    logging.debug(bcolors.FAIL + "Error with opening: " + opening_str)
+    tactics_file.close()
+    engine.quit()
+    sys.exit("INVALID/ILLEGAL OPENING")
+
+puzzle_count = 0
+
 while True:
     game = chess.pgn.read_game(all_games)
     if game is None:
         break
+    if opening != "NONE":
+        while not filter_game(opening, game):
+            game = chess.pgn.read_game(all_games)
+        if game is None:
+            break
+
     node = game
 
     game_id = game_id + 1
@@ -69,6 +112,10 @@ def prepare_settings():
     logging.debug(bcolors.OKGREEN + "Game Length: " + str(game.end().board().fullmove_number))
     logging.debug("Analysing Game..." + bcolors.ENDC)
 
+    # store last ply where black and white created a puzzle to avoid repeats
+    last_black = -1
+    last_white = -1
+
     while not node.is_end():
         next_node = node.variation(0)
 
@@ -79,9 +126,12 @@ def prepare_settings():
         logging.debug(bcolors.OKBLUE + "   CP: " + str(cur_score.score()))
         logging.debug("   Mate: " + str(cur_score.mate()) + bcolors.ENDC)
 
-        if investigate(prev_score, cur_score, node.board()):
-            logging.debug(bcolors.WARNING + "   Investigate!" + bcolors.ENDC)
-            puzzles.append(puzzle(node.board(), next_node.move, str(game_id), engine, info, game, settings.strict))
+        if min_turn * 2 - 1 <= node.ply() <= max_turn * 2 and (preferred_color == 2 or node.turn() == preferred_color):
+            if investigate(prev_score, cur_score, node.board(), checkmate_only):
+                last_colored_puzzle = last_white if node.ply()%2==1 else last_black
+                if node.ply() > last_colored_puzzle + 2:
+                    logging.debug(bcolors.WARNING + "   Investigate!" + bcolors.ENDC)
+                    puzzles.append(puzzle(node.board(), next_node.move, str(game_id), engine, info, game, settings.strict))
 
         prev_score = cur_score
         node = next_node
@@ -93,6 +143,9 @@ def prepare_settings():
             puzzle_pgn = post_puzzle(i, settings.include_blunder)
             tactics_file.write(puzzle_pgn)
             tactics_file.write("\n\n")
+            puzzle_count+=1
+
+logging.debug(bcolors.OKGREEN + "Puzzles generated:  " + str(puzzle_count))
 
 tactics_file.close()
 

modules/investigate/investigate.py

@@ -1,5 +1,5 @@
 import chess
-from chess import Board
+from chess import Board, pgn
 from chess.engine import Score
 
 
@@ -16,15 +16,14 @@ def material_count(board):
     return chess.popcount(board.occupied)
 
 
-def investigate(a: Score, b: Score, board: Board):
+def investigate(a: Score, b: Score, board: Board, checkmate_only: bool):
     """
     determine if the difference between position A and B
     is worth investigating for a puzzle.
     """
     a_cp, a_mate = a.score(), a.mate()
     b_cp, b_mate = b.score(), b.mate()
-
-    if a_cp is not None and b_cp is not None:
+    if a_cp is not None and b_cp is not None and not checkmate_only:
         if (((-110 < a_cp < 850 and 200 < b_cp < 850)
              or (-850 < a_cp < 110 and -200 > b_cp > -850))
                 and material_value(board) > 3
@@ -38,3 +37,16 @@ def investigate(a: Score, b: Score, board: Board):
         if sign(a_mate) == sign(b_mate):  # actually means that they're opposite
             return True
     return False
+
+
+def filter_game(opening: pgn, game_pgn: pgn):
+    if game_pgn is None or opening is None:
+        return True
+
+    moveIndex = 0
+
+    for (move, game_move) in zip(opening.mainline_moves(), game_pgn.mainline_moves()):
+        if game_move is None or game_move != move:
+            return False
+
+    return True

positions_for_investigation.py

@@ -41,6 +41,11 @@ def prepare_settings():
                         help="If False then generated puzzles won't include initial blunder move")
     parser.add_argument("--stockfish", metavar="STOCKFISH", default=None, help="Path to Stockfish binary")
 
+    parser.add_argument("--color", metavar="DESIRED_COLOR", type=int, default = 2, help="0 = black, 1 = white, 2 = any")
+
+    parser.add_argument("--forcedMate", metavar="MATES_ONLY", type=bool, default=False,
+                        help="only return forced checkmates")
+
     return parser.parse_args()
 
 

test/unit/test_regression.py

@@ -42,7 +42,7 @@ def test_investigate(self):
             print(score_b)
             logging.debug(f"Testing position {board.fen()} with scores {score_a} and {score_b}")
 
-            result = investigate(score_a, score_b, board)
+            result = investigate(score_a, score_b, board, False)
 
             self.assertEqual(expected_result, result)
 
@@ -97,6 +97,5 @@ def test_is_complete(self):
             logging.debug(f'{expected_result} vs {result}')
             self.assertEqual(expected_result, result)
 
-
 if __name__ == '__main__':
     unittest.main()