terrassumptions/server/app/gen/__init__.py

from typing import Dict, Tuple, List

from .shared import PointSource
from .map_crunch import MapCrunchPointSource
from .random_street_view import RSVPointSource
from .urban_centers import UrbanPointSource

from ..schemas import GenMethodEnum, GameConfig, CacheInfo

stock_target = 20

"""
Dictionary of PointSources
Maps (generation_method, only_america) -> PointSource
"""
sources: Dict[Tuple[GenMethodEnum, bool], PointSource] = {
    (GenMethodEnum.map_crunch, False): MapCrunchPointSource(stock_target=stock_target, max_retries=1000, only_america=False),
    (GenMethodEnum.map_crunch, True): MapCrunchPointSource(stock_target=stock_target, max_retries=1000, only_america=True),
    (GenMethodEnum.rsv, False): RSVPointSource(stock_target=stock_target, only_america=False),
    (GenMethodEnum.rsv, True): RSVPointSource(stock_target=stock_target, only_america=True),
    (GenMethodEnum.urban, False): UrbanPointSource(
        stock_target=stock_target,
        max_retries=100,
        retries_per_point=30,
        max_dist_km=25,
        usa_chance=0.1
    ),
    (GenMethodEnum.urban, True): UrbanPointSource(
        stock_target=stock_target,
        max_retries=100,
        retries_per_point=30,
        max_dist_km=25,
        usa_chance=1.0
    )
}


def generate_points(config: GameConfig) -> List[Tuple[float, float]]:
    """
    Generate points according to the GameConfig.
    """
    return sources[(config.generation_method, config.only_america)].get_points(config.rounds)


def restock_source(config: GameConfig):
    """
    Restock the PointSource associated with the GameConfig
    """
    sources[(config.generation_method, config.only_america)].restock()


def get_cache_info() -> List[CacheInfo]:
    """
    Get CacheInfo for each of the PointSources
    """
    return [CacheInfo(generation_method=gm, only_america=oa, size=len(src.stock)) for ((gm, oa), src) in sources.items()]