diff --git a/lib/engine_wrapper.py b/lib/engine_wrapper.py index cd89e725a..59a641fcb 100644 --- a/lib/engine_wrapper.py +++ b/lib/engine_wrapper.py @@ -18,6 +18,7 @@ from lib import config, model, lichess from lib.config import Configuration from lib.timer import Timer, msec, seconds, msec_str, sec_str, to_seconds +from lib.types import ReadableType from extra_game_handlers import game_specific_options from typing import Any, Optional, Union, Literal, Type from types import TracebackType @@ -367,18 +368,15 @@ def readable_number(self, number: int) -> str: def to_readable_value(self, stat: str, info: MOVE_INFO_TYPE) -> str: """Change a value to a more human-readable format.""" - readable: dict[str, Callable[[Any], str]] = {"Evaluation": self.readable_score, "Winrate": self.readable_wdl, - "Hashfull": lambda x: f"{round(x / 10, 1)}%", - "Nodes": self.readable_number, - "Speed": lambda x: f"{self.readable_number(x)}nps", - "Tbhits": self.readable_number, - "Cpuload": lambda x: f"{round(x / 10, 1)}%", - "Movetime": self.readable_time} - - def identity(x: Any) -> str: - return str(x) - - return str(readable.get(stat, identity)(info[stat])) + readable: ReadableType = {"Evaluation": self.readable_score, "Winrate": self.readable_wdl, + "Hashfull": lambda x: f"{round(x / 10, 1)}%", "Nodes": self.readable_number, + "Speed": lambda x: f"{self.readable_number(x)}nps", "Tbhits": self.readable_number, + "Cpuload": lambda x: f"{round(x / 10, 1)}%", "Movetime": self.readable_time} + + if stat in readable: + return readable[stat](info[stat]) + + return str(info[stat]) def get_stats(self, for_chat: bool = False) -> list[str]: """ diff --git a/lib/lichess.py b/lib/lichess.py index e18d44357..5d1526bee 100644 --- a/lib/lichess.py +++ b/lib/lichess.py @@ -12,8 +12,8 @@ from lib.timer import Timer, seconds, sec_str from typing import Optional, Union, Any import chess.engine -JSON_REPLY_TYPE = dict[str, Any] -REQUESTS_PAYLOAD_TYPE = dict[str, Any] +from lib.types import UserProfileType, JSON_REPLY_TYPE, REQUESTS_PAYLOAD_TYPE + ENDPOINTS = { "profile": "/api/account", @@ -322,9 +322,9 @@ def decline_challenge(self, challenge_id: str, reason: str = "generic") -> None: except Exception: pass - def get_profile(self) -> JSON_REPLY_TYPE: + def get_profile(self) -> UserProfileType: """Get the bot's profile (e.g. username).""" - profile = self.api_get_json("profile") + profile: UserProfileType = self.api_get_json("profile") self.set_user_agent(profile["username"]) return profile @@ -353,7 +353,7 @@ def get_game_pgn(self, game_id: str) -> str: except Exception: return "" - def get_online_bots(self) -> list[dict[str, Any]]: + def get_online_bots(self) -> list[UserProfileType]: """Get a list of bots that are online.""" try: online_bots_str = self.api_get_raw("online_bots") diff --git a/lib/matchmaking.py b/lib/matchmaking.py index f69f8cd4a..1d32dfcdf 100644 --- a/lib/matchmaking.py +++ b/lib/matchmaking.py @@ -10,7 +10,8 @@ from lib import lichess from lib.config import Configuration, FilterType from typing import Any, Optional, Union -USER_PROFILE_TYPE = dict[str, Any] +from lib.types import UserProfileType, PerfType +USER_PROFILE_TYPE = UserProfileType EVENT_TYPE = dict[str, Any] MULTIPROCESSING_LIST_TYPE = Sequence[model.Challenge] DAILY_TIMERS_TYPE = list[Timer] @@ -132,9 +133,9 @@ def update_daily_challenge_record(self) -> None: self.min_wait_time = seconds(60) * ((len(self.daily_challenges) // 50) + 1) write_daily_challenges(self.daily_challenges) - def perf(self) -> dict[str, dict[str, Any]]: + def perf(self) -> dict[str, PerfType]: """Get the bot's rating in every variant. Bullet, blitz, rapid etc. are considered different variants.""" - user_perf: dict[str, dict[str, Any]] = self.user_profile["perfs"] + user_perf: dict[str, PerfType] = self.user_profile["perfs"] return user_perf def username(self) -> str: @@ -155,7 +156,9 @@ def get_weights(self, online_bots: list[USER_PROFILE_TYPE], rating_preference: s game_type: str) -> list[int]: """Get the weight for each bot. A higher weights means the bot is more likely to get challenged.""" def rating(bot: USER_PROFILE_TYPE) -> int: - return int(bot.get("perfs", {}).get(game_type, {}).get("rating", 0)) + perfs: dict[str, PerfType] = bot.get("perfs", {}) + perf: PerfType = perfs.get(game_type, {}) + return perf.get("rating", 0) if rating_preference == "high": # A bot with max_rating rating will be twice as likely to get picked than a bot with min_rating rating. diff --git a/lib/model.py b/lib/model.py index 72091db17..1b240bb82 100644 --- a/lib/model.py +++ b/lib/model.py @@ -8,6 +8,7 @@ from lib.config import Configuration from typing import Any from collections import defaultdict +from lib.types import UserProfileType logger = logging.getLogger(__name__) @@ -15,7 +16,7 @@ class Challenge: """Store information about a challenge.""" - def __init__(self, challenge_info: dict[str, Any], user_profile: dict[str, Any]) -> None: + def __init__(self, challenge_info: dict[str, Any], user_profile: UserProfileType) -> None: """:param user_profile: Information about our bot.""" self.id = challenge_info["id"] self.rated = challenge_info["rated"] diff --git a/lib/types.py b/lib/types.py new file mode 100644 index 000000000..90a38680b --- /dev/null +++ b/lib/types.py @@ -0,0 +1,59 @@ +from typing import NotRequired, TypedDict, Any, Callable +from chess.engine import PovWdl, PovScore + +JSON_REPLY_TYPE = dict[str, Any] +REQUESTS_PAYLOAD_TYPE = dict[str, Any] + + +class PerfType(TypedDict): + games: NotRequired[int] + rating: NotRequired[int] + rd: NotRequired[int] + sd: NotRequired[int] + prov: NotRequired[bool] + + +class ProfileType(TypedDict): + country: NotRequired[str] + location: NotRequired[str] + bio: NotRequired[str] + firstName: NotRequired[str] + lastName: NotRequired[str] + fideRating: NotRequired[int] + uscfRating: NotRequired[int] + ecfRating: NotRequired[int] + cfcRating: NotRequired[int] + dsbRating: NotRequired[int] + links: NotRequired[str] + + +class UserProfileType(TypedDict): + id: str + username: str + perfs: NotRequired[dict[str, PerfType]] + createdAt: NotRequired[int] + disabled: NotRequired[bool] + tosViolation: NotRequired[bool] + profile: NotRequired[ProfileType] + seenAt: NotRequired[int] + patron: NotRequired[int] + verified: NotRequired[int] + playTime: NotRequired[dict[str, int]] + title: NotRequired[str] + online: NotRequired[bool] + url: NotRequired[str] + followable: NotRequired[bool] + following: NotRequired[bool] + blocking: NotRequired[bool] + followsYou: NotRequired[bool] + + +class ReadableType(TypedDict): + Evaluation: Callable[[PovScore], str] + Winrate: Callable[[PovWdl], str] + Hashfull: Callable[[int], str] + Nodes: Callable[[int], str] + Speed: Callable[[int], str] + Tbhits: Callable[[int], str] + Cpuload: Callable[[int], str] + Movetime: Callable[[int], str] diff --git a/lichess-bot.py b/lichess-bot.py index 8955ead1e..36b795bd4 100644 --- a/lichess-bot.py +++ b/lichess-bot.py @@ -25,6 +25,7 @@ from lib.config import load_config, Configuration from lib.conversation import Conversation, ChatLine from lib.timer import Timer, seconds, msec, hours, to_seconds +from lib.types import UserProfileType from requests.exceptions import ChunkedEncodingError, ConnectionError, HTTPError, ReadTimeout from rich.logging import RichHandler from collections import defaultdict @@ -33,7 +34,8 @@ from queue import Queue, Empty from multiprocessing.pool import Pool from typing import Any, Optional, Union -USER_PROFILE_TYPE = dict[str, Any] +from types import FrameType +USER_PROFILE_TYPE = UserProfileType EVENT_TYPE = dict[str, Any] PLAY_GAME_ARGS_TYPE = dict[str, Any] EVENT_GETATTR_GAME_TYPE = dict[str, Any] @@ -63,7 +65,7 @@ def disable_restart() -> None: restart = False -def signal_handler(signal: int, frame: Any) -> None: +def signal_handler(signal: int, frame: Optional[FrameType]) -> None: """Terminate lichess-bot.""" global terminated global force_quit diff --git a/test_bot/lichess.py b/test_bot/lichess.py index e85e52096..bd7dcfacb 100644 --- a/test_bot/lichess.py +++ b/test_bot/lichess.py @@ -7,9 +7,10 @@ import traceback import datetime from queue import Queue -from typing import Union, Any, Optional, Generator +from typing import Any, Optional, Generator from lib.timer import to_msec -from lib.lichess import JSON_REPLY_TYPE, REQUESTS_PAYLOAD_TYPE +from lib.types import UserProfileType, JSON_REPLY_TYPE, REQUESTS_PAYLOAD_TYPE + logger = logging.getLogger(__name__) @@ -187,7 +188,7 @@ def decline_challenge(self, challenge_id: str, reason: str = "generic") -> None: """Isn't used in tests.""" pass - def get_profile(self) -> dict[str, Union[str, bool, dict[str, str]]]: + def get_profile(self) -> UserProfileType: """Return a simple profile for the bot that lichess-bot uses when testing.""" return {"id": "b", "username": "b", @@ -222,9 +223,9 @@ def get_game_pgn(self, game_id: str) -> str: * """ - def get_online_bots(self) -> list[dict[str, Union[str, bool]]]: + def get_online_bots(self) -> list[UserProfileType]: """Return that the only bot online is us.""" - return [{"username": "b", "online": True}] + return [{"username": "b", "id": "b", "online": True}] def challenge(self, username: str, payload: REQUESTS_PAYLOAD_TYPE) -> JSON_REPLY_TYPE: """Isn't used in tests."""