version 1.2.2: WGApiWoTBlitzTankopedia(): add tier cache for faster lookups by tier

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "blitz-models"
-version = "1.2.1"
+version = "1.2.2"
 authors = [{ name = "Jylpah", email = "[email protected]" }]
 description = "Pydantic models for Wargaming's World of Tanks Blitz game "
 readme = "README.md"

src/blitzmodels/__init__.py

@@ -1,5 +1,5 @@
 from .config import get_config_file as get_config_file
-
+from .types import AccountId as AccountId, TankId as TankId
 from .region import Region as Region
 from .release import Release as Release
 from .account import Account as Account
@@ -38,6 +38,7 @@
 
 
 __all__ = [
+    "types",
     "account",
     "config",
     "map",

src/blitzmodels/account.py

@@ -19,6 +19,7 @@
 
 from .region import Region
 from .wg_api import AccountInfo
+from .types import AccountId
 
 logger = logging.getLogger()
 error = logger.error
@@ -37,15 +38,14 @@
 
 TypeAccountDict = dict[str, int | bool | Region | None]
 
-
 # def lateinit_region() -> Region:
 #     """Required for initializing a model w/o a 'region' field"""
 #     raise RuntimeError("lateinit_region(): should never be called")
 
 
 class Account(JSONExportable, CSVExportable, TXTExportable, TXTImportable, Importable):
     # fmt: off
-    id              : int       = Field(alias="_id")
+    id              : AccountId = Field(alias="_id")
     # lateinit is a trick to fool mypy since region is set in root_validator
     region          : Region    = Field(default=Region.bot, alias="r")
     last_battle_time: int       = Field(default=0, alias="l")

src/blitzmodels/tank.py

@@ -14,6 +14,8 @@
     TEXT,
 )
 
+from .types import TankId
+
 logger = logging.getLogger()
 error = logger.error
 message = logger.warning
@@ -119,7 +121,7 @@ def __str__(self) -> str:
 
 class Tank(JSONExportable, CSVExportable, TXTExportable):
     # fmt: off
-    tank_id 	: int 				= Field(default=..., alias = '_id')
+    tank_id 	: TankId			= Field(default=..., alias = '_id')
     name   		: str 				= Field(default="")
     code        : str | None        = Field(default=None)
     nation   	: EnumNation  		= Field(default=EnumNation.european)

src/blitzmodels/types.py

@@ -0,0 +1,4 @@
+## Type aliases
+
+AccountId = int
+TankId = int

src/blitzmodels/wg_api.py

@@ -5,14 +5,14 @@
     TypeVar,
     Sequence,
     Tuple,
+    Set,
     Self,
     Type,
     Dict,
     Annotated,
 )
 from types import TracebackType
 import logging
-from sys import path
 import pyarrow  # type: ignore
 from bson import ObjectId
 from pydantic import (
@@ -46,18 +46,15 @@
 from pyutils.utils import epoch_now
 from pyutils import ThrottledClientSession
 
-# Fix relative imports
-from pathlib import Path
-
-path.insert(0, str(Path(__file__).parent.parent.resolve()))
-
-from blitzmodels.region import Region  # noqa: E402
-from blitzmodels.tank import (  # noqa: E402
+from .region import Region
+from .tank import (
     Tank,
     EnumNation,
     EnumVehicleTypeStr,
     EnumVehicleTier,
 )
+from .types import AccountId, TankId
+
 
 TYPE_CHECKING = True
 logger = logging.getLogger()
@@ -83,8 +80,7 @@
     message: str | None
     field: str | None
     value: str | None
-    # TODO[pydantic]: The following keys were removed: `allow_mutation`.
-    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+
     model_config = ConfigDict(
         frozen=False, validate_assignment=True, populate_by_name=True
     )
@@ -93,8 +89,6 @@
         return f"code: {self.code} {self.message}"
 
 
-
-
 class WGTankStatAll(JSONExportable):
     # fmt: off
     battles:            int = Field(..., alias="b")
@@ -132,8 +126,8 @@
 
 
 class AccountInfoStats(WGTankStatAll):
-    max_frags_tank_id   : int = Field(default=0, alias="mft")
-    max_xp_tank_id      : int = Field(default=0, alias="mxt")
+    max_frags_tank_id: int = Field(default=0, alias="mft")
+    max_xp_tank_id: int = Field(default=0, alias="mxt")
 
 
 class TankStat(JSONExportable):
@@ -142,8 +136,8 @@
     region:             Region | None = Field(default=None, alias="r")
     all:                WGTankStatAll = Field(..., alias="s")
     last_battle_time:   int = Field(..., alias="lb")
-    account_id:         int = Field(..., alias="a")
-    tank_id:            int = Field(..., alias="t")
+    account_id:         TankId = Field(..., alias="a")
+    tank_id:            TankId = Field(..., alias="t")
     mark_of_mastery:    int = Field(default=0, alias="m")
     battle_life_time:   int = Field(default=0, alias="l")
     release:            str | None = Field(default=None, alias="u")
@@ -274,7 +268,7 @@
 
     @classmethod
     def mk_id(
-        cls, account_id: int, last_battle_time: int, tank_id: int = 0
+        cls, account_id: AccountId, last_battle_time: int, tank_id: TankId = 0
     ) -> ObjectId:
         return ObjectId(
             hex(account_id)[2:].zfill(10)
@@ -317,14 +311,14 @@
                     tank_id={self.tank_id} \
                     last_battle_time={self.last_battle_time}"
 
+
 ###########################################
 #
 # AccountInfo()
 #
 ###########################################
 
 
-
 class AccountInfo(JSONExportable):
     # fmt: off
     account_id:         int = Field(alias="id")
@@ -337,7 +331,7 @@
     # fmt: on
 
     model_config = ConfigDict(
-        # arbitrary_types_allowed=True,  # should this be removed? 
+        # arbitrary_types_allowed=True,  # should this be removed?
         frozen=False,
         validate_assignment=True,
         populate_by_name=True,
@@ -405,6 +399,7 @@
                 "nickname": "jylpah"
             }"""
 
+
 class WGApiWoTBlitz(JSONExportable):
     # fmt: off
     status: str = Field(default="ok", alias="s")
@@ -594,8 +589,7 @@
     max_series: PlayerAchievementsMaxSeries | None = Field(default=None, alias="m")
     account_id: int | None = Field(default=None)
     updated: int | None = Field(default=None)
-    # TODO[pydantic]: The following keys were removed: `allow_mutation`.
-    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+
     model_config = ConfigDict(
         frozen=False, validate_assignment=True, populate_by_name=True
     )
@@ -609,8 +603,7 @@
 
 class WGApiWoTBlitzPlayerAchievements(WGApiWoTBlitz):
     data: dict[str, PlayerAchievementsMain] | None = Field(default=None, alias="d")
-    # TODO[pydantic]: The following keys were removed: `allow_mutation`.
-    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+
     model_config = ConfigDict(
         frozen=False, validate_assignment=True, populate_by_name=True
     )
@@ -674,8 +667,12 @@
     data: Dict[str, Tank] = Field(default=dict(), alias="d")
     codes: Dict[str, Tank] = Field(default=dict(), alias="c")
 
+    # TODO: Implement tier cache
+    _tier_cache: Dict[int, Set[TankId]] = dict()
+
     _exclude_export_DB_fields = {"codes": True}
     _exclude_export_src_fields = {"codes": True}
+
     model_config = ConfigDict(
         frozen=False, validate_assignment=True, populate_by_name=True
     )
@@ -684,12 +681,14 @@
     def _validate_code(self) -> Self:
         if len(self.codes) == 0:
             self._set_skip_validation("codes", self._update_codes(data=self.data))
+        if len(self._tier_cache) == 0:
+            self._set_skip_validation("_tier_cache", self._update_tier_cache())
         return self
 
     def __len__(self) -> int:
         return len(self.data)
 
-    def __getitem__(self, key: str | int) -> Tank:
+    def __getitem__(self, key: str | TankId) -> Tank:
         if isinstance(key, int):
             key = str(key)
         return self.data[key]
@@ -703,6 +702,15 @@
             self.meta = dict()
         self.meta["count"] = len(self.data)
 
+    def _update_tier_cache(self) -> Dict[int, Set[TankId]]:
+        """Update tier cache and return new cache"""
+        res: Dict[int, Set[TankId]] = dict()
+        for tier in range(1, 11):
+            res[tier] = set()
+        for tank in self.data.values():
+            res[tank.tier].add(tank.tank_id)
+        return res
+
     def _code_add(self, tank: Tank, codes: dict[str, Tank]) -> bool:
         if tank.code is not None:
             codes[tank.code] = tank
@@ -711,16 +719,18 @@
 
     def add(self, tank: Tank) -> None:
         self.data[str(tank.tank_id)] = tank
+        self._tier_cache[tank.tier].add(tank.tank_id)
         self._code_add(tank, self.codes)
         self.update_count()
 
-    def pop(self, tank_id: int) -> Tank:
+    def pop(self, tank_id: TankId) -> Tank:
         """Raises KeyError if tank_id is not found in self.data"""
         tank: Tank = self.data.pop(str(tank_id))
         self.update_count()
         if tank.code is not None:
             try:
                 del self.codes[tank.code]
+                self._tier_cache[tank.tier].remove(tank.tank_id)
             except Exception as err:
                 debug(f"could not remove code for tank_id={tank.tank_id}: {err}")
                 pass
@@ -751,23 +761,30 @@
         """update _code dict"""
         self._set_skip_validation("codes", self._update_codes(self.data))
 
-    def update_tanks(self, new: "WGApiWoTBlitzTankopedia") -> Tuple[set[int], set[int]]:
+    def update_tanks(
+        self, new: "WGApiWoTBlitzTankopedia"
+    ) -> Tuple[set[TankId], set[TankId]]:
         """update tankopedia with another one"""
-        new_ids: set[int] = {tank.tank_id for tank in new}
-        old_ids: set[int] = {tank.tank_id for tank in self}
-        added: set[int] = new_ids - old_ids
-        updated: set[int] = new_ids & old_ids
+        new_ids: set[TankId] = {tank.tank_id for tank in new}
+        old_ids: set[TankId] = {tank.tank_id for tank in self}
+        added: set[TankId] = new_ids - old_ids
+        updated: set[TankId] = new_ids & old_ids
         updated = {tank_id for tank_id in updated if new[tank_id] != self[tank_id]}
 
         self.data.update({(str(tank_id), new[tank_id]) for tank_id in added})
-        updated_ids: set[int] = set()
+        updated_ids: set[TankId] = set()
         for tank_id in updated:
             if self.data[str(tank_id)].update(new[tank_id]):
                 updated_ids.add(tank_id)
         self.update_count()
         self.update_codes()
         return (added, updated_ids)
 
+    def get_tank_ids_by_tier(self, tier: int) -> Set[TankId]:
+        if tier < 1 or tier > 10:
+            raise ValueError(f"tier must be between 1-10: {tier}")
+        return self._tier_cache[tier]
+
 
 class WGApiTankString(JSONExportable):
     id: int

tests/test_tank.py

@@ -249,6 +249,12 @@ async def test_10_WGApiTankopedia(
                 assert (
                     False
                 ), f"Parsing test file List[Tank] failed: {basename(tanks_json_fn)}"
+    N: int = 0
+    for tier in range(1, 11):
+        N += len(tankopedia.get_tank_ids_by_tier(tier=tier))
+    assert N == len(
+        tankopedia
+    ), f"incorrect number of tanks in the tier cache: {N} != {len(tankopedia)}"
     for tank in tanks_json:
         tankopedia.add(tank)
     debug("read %d tanks", len(tankopedia.data))
@@ -296,9 +302,7 @@ async def test_11_WGApiTankopedia(
                     await file.read()
                 )
             except Exception:
-                assert (
-                    False
-                ), f"Parsing test file WGApiWoTBlitzTankopedia() failed: {basename(tankopedia_fn)}"
+                assert False, f"Parsing test file WGApiWoTBlitzTankopedia() failed: {basename(tankopedia_fn)}"
 
     debug("read %d tanks", len(tankopedia.data))
     assert tankopedia.meta is not None, "Failed to update meta"
@@ -309,9 +313,7 @@ async def test_11_WGApiTankopedia(
         len(tankopedia.data) == tankopedia_tanks
     ), f"could not import all the tanks: got {tankopedia.data}, should be {tankopedia_tanks}"
 
-    assert (
-        tankopedia.has_codes
-    ), f"could not generate all the codes: tanks={len(tankopedia.data)}, codes={len(tankopedia.codes)}"
+    assert tankopedia.has_codes, f"could not generate all the codes: tanks={len(tankopedia.data)}, codes={len(tankopedia.codes)}"
     # test tankopedia export import
     tankopedia_file: str = f"{tmp_path.resolve()}/tankopedia.json"
     try:
@@ -347,9 +349,7 @@ async def test_12_WGApiTankopedia_sorted(
                     await file.read()
                 )
             except Exception:
-                assert (
-                    False
-                ), f"Parsing test file WGApiWoTBlitzTankopedia() failed: {basename(tankopedia_fn)}"
+                assert False, f"Parsing test file WGApiWoTBlitzTankopedia() failed: {basename(tankopedia_fn)}"
 
     debug("read %d tanks", len(tankopedia.data))
     tanks: list[Tank] = list()