diff --git a/game/lasercodes/__init__.py b/game/lasercodes/__init__.py new file mode 100644 index 00000000..4b2e1986 --- /dev/null +++ b/game/lasercodes/__init__.py @@ -0,0 +1 @@ +from .lasercoderegistry import LaserCodeRegistry diff --git a/game/missiongenerator/lasercoderegistry.py b/game/lasercodes/lasercoderegistry.py similarity index 53% rename from game/missiongenerator/lasercoderegistry.py rename to game/lasercodes/lasercoderegistry.py index 6872cb30..2b4226b3 100644 --- a/game/missiongenerator/lasercoderegistry.py +++ b/game/lasercodes/lasercoderegistry.py @@ -1,30 +1,21 @@ from collections import deque -from typing import Iterator - - -class OutOfLaserCodesError(RuntimeError): - def __init__(self) -> None: - super().__init__( - f"All JTAC laser codes have been allocated. No available codes." - ) class LaserCodeRegistry: def __init__(self) -> None: self.allocated_codes: set[int] = set() - self.allocator: Iterator[int] = LaserCodeRegistry.__laser_code_generator() + self.available_codes = LaserCodeRegistry._all_valid_laser_codes() - def get_next_laser_code(self) -> int: + def alloc_laser_code(self) -> int: try: - while (code := next(self.allocator)) in self.allocated_codes: - pass + code = self.available_codes.popleft() self.allocated_codes.add(code) return code - except StopIteration: - raise OutOfLaserCodesError + except IndexError: + raise RuntimeError("All laser codes have been allocated") @staticmethod - def __laser_code_generator() -> Iterator[int]: + def _all_valid_laser_codes() -> deque[int]: # Valid laser codes are as follows # First digit is always 1 # Second digit is 5-7 @@ -34,4 +25,4 @@ class LaserCodeRegistry: # We start with the default of 1688 and wrap around when we reach the end q.rotate(-q.index(1688)) - return iter(q) + return q diff --git a/game/missiongenerator/aircraft/aircraftgenerator.py b/game/missiongenerator/aircraft/aircraftgenerator.py index fcbc5121..37693ffb 100644 --- a/game/missiongenerator/aircraft/aircraftgenerator.py +++ b/game/missiongenerator/aircraft/aircraftgenerator.py @@ -18,7 +18,7 @@ from game.ato.package import Package from game.ato.starttype import StartType from game.factions.faction import Faction from game.missiongenerator.missiondata import MissionData -from game.missiongenerator.lasercoderegistry import LaserCodeRegistry +from game.lasercodes import LaserCodeRegistry from game.radio.radios import RadioRegistry from game.radio.tacan import TacanRegistry from game.runways import RunwayData diff --git a/game/missiongenerator/aircraft/flightgroupconfigurator.py b/game/missiongenerator/aircraft/flightgroupconfigurator.py index 437c31c1..97530488 100644 --- a/game/missiongenerator/aircraft/flightgroupconfigurator.py +++ b/game/missiongenerator/aircraft/flightgroupconfigurator.py @@ -12,7 +12,7 @@ from dcs.unitgroup import FlyingGroup from game.ato import Flight, FlightType from game.callsigns import callsign_for_support_unit from game.data.weapons import Pylon, WeaponType as WeaponTypeEnum -from game.missiongenerator.lasercoderegistry import LaserCodeRegistry +from game.lasercodes import LaserCodeRegistry from game.missiongenerator.logisticsgenerator import LogisticsGenerator from game.missiongenerator.missiondata import AwacsInfo, MissionData, TankerInfo from game.radio.radios import RadioFrequency, RadioRegistry @@ -133,7 +133,7 @@ class FlightGroupConfigurator: ) -> None: self.set_skill(unit, member) if member.loadout.has_weapon_of_type(WeaponTypeEnum.TGP) and member.is_player: - laser_codes.append(self.laser_code_registry.get_next_laser_code()) + laser_codes.append(self.laser_code_registry.alloc_laser_code()) else: laser_codes.append(None) diff --git a/game/missiongenerator/flotgenerator.py b/game/missiongenerator/flotgenerator.py index af7a97a3..3c5e960c 100644 --- a/game/missiongenerator/flotgenerator.py +++ b/game/missiongenerator/flotgenerator.py @@ -43,7 +43,7 @@ from game.theater.controlpoint import ControlPoint from game.unitmap import UnitMap from game.utils import Heading from .frontlineconflictdescription import FrontLineConflictDescription -from .lasercoderegistry import LaserCodeRegistry +from game.lasercodes import LaserCodeRegistry from .missiondata import JtacInfo, MissionData if TYPE_CHECKING: @@ -144,7 +144,7 @@ class FlotGenerator: if self.game.lua_plugin_manager.is_option_enabled("ctld", "fc3LaserCode"): code = 1113 else: - code = self.laser_code_registry.get_next_laser_code() + code = self.laser_code_registry.alloc_laser_code() utype = self.game.blue.faction.jtac_unit if utype is None: diff --git a/game/missiongenerator/missiongenerator.py b/game/missiongenerator/missiongenerator.py index 30fd0c54..b4e932e8 100644 --- a/game/missiongenerator/missiongenerator.py +++ b/game/missiongenerator/missiongenerator.py @@ -34,7 +34,7 @@ from .flotgenerator import FlotGenerator from .forcedoptionsgenerator import ForcedOptionsGenerator from .frontlineconflictdescription import FrontLineConflictDescription from .kneeboard import KneeboardGenerator -from .lasercoderegistry import LaserCodeRegistry +from game.lasercodes import LaserCodeRegistry from .luagenerator import LuaGenerator from .missiondata import MissionData from .tgogenerator import TgoGenerator diff --git a/tests/lasercodes/__init__.py b/tests/lasercodes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/lasercodes/test_lasercoderegistry.py b/tests/lasercodes/test_lasercoderegistry.py new file mode 100644 index 00000000..be5d595d --- /dev/null +++ b/tests/lasercodes/test_lasercoderegistry.py @@ -0,0 +1,15 @@ +from game.lasercodes.lasercoderegistry import LaserCodeRegistry + + +def test_initial_laser_codes() -> None: + reg = LaserCodeRegistry() + assert list(reg.available_codes)[:5] == [1688, 1687, 1686, 1685, 1684] + assert list(reg.available_codes)[-5:] == [1715, 1714, 1713, 1712, 1711] + assert len(reg.available_codes) == 192 + + +def test_alloc_laser_code() -> None: + reg = LaserCodeRegistry() + assert reg.alloc_laser_code() == 1688 + assert 1688 not in reg.available_codes + assert len(reg.available_codes) == 191