Allow the user to optionally specify a movetime for get_wdl_stats. (zhelyabuzhsky#139)

* Allow the user to optionally specify a movetime for `get_wdl_stats`.

* Explain the arguments below the example.

Author: johndoknjas

README.md

@@ -265,13 +265,17 @@ Before calling this function, it is recommended that you first check if your ver
 use the "does_current_engine_version_have_wdl_option()" function below.
 
 ```python
-stockfish.get_wdl_stats() # include the argument get_as_tuple=True if you'd like to have a tuple returned instead of a list.
+stockfish.get_wdl_stats()
 ```
 
 ```text
 [87, 894, 19]
 ```
 
+Optional arguments:
+- `get_as_tuple`: if you'd like to have a tuple returned instead of a list.
+- `time`: if you'd like to constrain the search by a duration rather than the current depth.
+
 ### Find if your version of Stockfish is recent enough to display WDL stats
 
 ```python

stockfish/models.py

@@ -542,7 +542,7 @@ def get_best_move(
         self, wtime: int | None = None, btime: int | None = None
     ) -> str | None:
         """Returns the best move in the current position on the board.
-        `wtime` and `btime` arguments influence the search only if provided.
+        If both `wtime` and `btime` aren't provided, the current depth is used for the search.
 
         Args:
             wtime:
@@ -686,13 +686,16 @@ def is_move_correct(self, move_value: str) -> bool:
         return move_value in self.get_perft(1)[1]
 
     def get_wdl_stats(
-        self, get_as_tuple: bool = False
+        self, get_as_tuple: bool = False, time: int | None = None
     ) -> list[int] | tuple[int, int, int] | None:
         """Returns Stockfish's win/draw/loss stats for the side to move.
 
         Args:
             get_as_tuple:
                 Option to return the wdl stats as a tuple instead of a list. Default is `False`.
+            time:
+                Time for Stockfish to search (milliseconds). If provided, will be used instead of the
+                current depth.
 
         Returns:
             A list or tuple of three integers, unless the game is over (in which case
@@ -709,18 +712,17 @@ def get_wdl_stats(
                 + """ get_wdl_stats will still return full strength Stockfish's wdl stats of the position."""
             )
 
-        self._go()
+        if time is None:
+            self._go()
+        else:
+            self._go_time(time)
         lines = self._get_sf_go_command_output()
         if lines[-1].startswith("bestmove (none)"):
             return None
         split_line = [line.split(" ") for line in lines if " multipv 1 " in line][-1]
         wdl_index = split_line.index("wdl")
-
         wdl_stats = [int(split_line[i]) for i in range(wdl_index + 1, wdl_index + 4)]
-
-        if get_as_tuple:
-            return (wdl_stats[0], wdl_stats[1], wdl_stats[2])
-        return wdl_stats
+        return (wdl_stats[0], wdl_stats[1], wdl_stats[2]) if get_as_tuple else wdl_stats
 
     def does_current_engine_version_have_wdl_option(self) -> bool:
         """Returns whether the user's version of Stockfish has the option to display WDL stats."""

tests/stockfish/test_models.py

@@ -55,6 +55,7 @@ def test_get_best_move_first_move(self, stockfish: Stockfish):
         best_move = stockfish.get_best_move()
         assert best_move in ("e2e3", "e2e4", "g1f3", "b1c3", "d2d4")
 
+    @pytest.mark.slow
     def test_get_best_move_time_first_move(self, stockfish: Stockfish):
         best_move = stockfish.get_best_move_time(1000)
         assert best_move in ("e2e3", "e2e4", "g1f3", "b1c3", "d2d4")
@@ -83,7 +84,8 @@ def test_info_raises_error_by_default(self, stockfish: Stockfish):
         stockfish.get_evaluation()
         stockfish.get_perft(1)
         stockfish.get_top_moves(1)
-        stockfish.get_wdl_stats()
+        if stockfish.does_current_engine_version_have_wdl_option():
+            stockfish.get_wdl_stats()
         with pytest.raises(RuntimeError):
             stockfish.info()
         stockfish.get_best_move_time(1)
@@ -94,6 +96,7 @@ def test_get_best_move_not_first_move(self, stockfish: Stockfish):
         best_move = stockfish.get_best_move()
         assert best_move in ("d2d4", "g1f3")
 
+    @pytest.mark.slow
     def test_get_best_move_time_not_first_move(self, stockfish: Stockfish):
         stockfish.make_moves_from_start(["e2e4", "e7e6"])
         best_move = stockfish.get_best_move_time(1000)
@@ -935,6 +938,15 @@ def test_get_wdl_stats(self, stockfish: Stockfish):
             with pytest.raises(RuntimeError):
                 stockfish.get_wdl_stats()
 
+    @pytest.mark.slow
+    def test_get_wdl_stats_movetime(self, stockfish: Stockfish):
+        if not stockfish.does_current_engine_version_have_wdl_option():
+            return
+        start = time.time()
+        wdl_stats = stockfish.get_wdl_stats(get_as_tuple=True, time=1000)
+        assert 0.5 < time.time() - start < 1.5
+        assert wdl_stats and wdl_stats[1] > wdl_stats[0] > wdl_stats[2]
+
     def test_does_current_engine_version_have_wdl_option(self, stockfish: Stockfish):
         if stockfish.get_stockfish_major_version() <= 11:
             assert not stockfish.does_current_engine_version_have_wdl_option()