diff --git a/cp_red_npc_generator/src/generate_weapon.py b/cp_red_npc_generator/src/generate_weapon.py index 80f2dcf..52616d5 100644 --- a/cp_red_npc_generator/src/generate_weapon.py +++ b/cp_red_npc_generator/src/generate_weapon.py @@ -124,33 +124,12 @@ def get_brawling_weapon_item(npc: Npc, all_weapons: List[ItemWithNames]) -> Item name="Brawling") -def add_weapon_to_npc(weapon: Item, npc: Npc, weapon_skills_data: Dict[str, str]): - skill_name: Optional[str] = None - for tag in weapon.get_all_tags(): - if tag in weapon_skills_data.keys(): - skill_name = weapon_skills_data[tag] - break - - assert skill_name - skill_value: int = npc.get_skill_total_value(skill_name) - if weapon.quality == ItemQuality.EXCELLENT: - skill_value += 1 - - weapon = replace( - weapon, - name=f"[{skill_value}] {weapon.name}") - - npc.weapons.add(weapon) - - def generate_weapon(npc: Npc, npc_template: NpcTemplate) -> Npc: logging.debug("\nGenerating weapons...") weapons_data = load_data("configs/items/weapon.json") all_weapons: List[ItemWithNames] = [dataclass_wizard.fromdict(ItemWithNames, x) for x in weapons_data] - weapon_skills_data: Dict[str, str] = load_data("configs/weapon_skills.json") - total_weapons_budget: int = round(npc_template.rank.items_budget[ItemType.WEAPON].generate()) logging.debug(f"\t{total_weapons_budget=}") @@ -172,7 +151,7 @@ def generate_weapon(npc: Npc, npc_template: NpcTemplate) -> Npc: npc) if primary_weapon: - add_weapon_to_npc(primary_weapon, npc, weapon_skills_data) + npc.weapons.add(primary_weapon) # try to buy a secondary weapon with any budget left secondary_weapon, _ = pick_weapon( @@ -183,14 +162,14 @@ def generate_weapon(npc: Npc, npc_template: NpcTemplate) -> Npc: npc) if secondary_weapon: - add_weapon_to_npc(secondary_weapon, npc, weapon_skills_data) + npc.weapons.add(secondary_weapon) # add all the rest weapons from the cyberware for cyberware in npc.cyberware: if cyberware.item.damage: - add_weapon_to_npc(cyberware.item, npc, weapon_skills_data) + npc.weapons.add(cyberware.item) # add boxing or martial arts - add_weapon_to_npc(get_brawling_weapon_item(npc, all_weapons), npc, weapon_skills_data) + npc.weapons.add(get_brawling_weapon_item(npc, all_weapons)) return npc diff --git a/cp_red_npc_generator/src/npc.py b/cp_red_npc_generator/src/npc.py index 36d59c8..ae162ea 100644 --- a/cp_red_npc_generator/src/npc.py +++ b/cp_red_npc_generator/src/npc.py @@ -5,13 +5,14 @@ import logging import math from functools import cmp_to_key -from typing import Dict, List, Tuple, Set -from dataclasses import dataclass, field +from typing import Dict, List, Tuple, Set, Optional +from dataclasses import dataclass, field, replace +from cp_red_npc_generator.src.utils import load_data from npc_template import TraumaTeamStatusType from modifier import ModifierSource from inventory_node import InventoryNode -from item import Item, ItemType +from item import Item, ItemType, ItemQuality from stats import StatType, Skill, SkillType, StatSkillValue from table_view import TableView @@ -196,13 +197,33 @@ def armor_sorter(item: Item) -> int: def weapon_sorter(item: Item) -> int: return int(item.damage[0]) * int(item.damage[2]) * item.rate_of_fire - sorted_melee_weapon = sorted([x for x in self.weapons if "MeleeWeapon" in x.get_all_tags()], - key=weapon_sorter, - reverse=True) - - sorted_ranged_weapon = sorted([x for x in self.weapons if "RangedWeapon" in x.get_all_tags()], - key=weapon_sorter, - reverse=True) + weapon_skills_data: Dict[str, str] = load_data("configs/weapon_skills.json") + + def add_skill_value(weapon: Item) -> Item: + skill_name: Optional[str] = None + for tag in weapon.get_all_tags(): + if tag in weapon_skills_data.keys(): + skill_name = weapon_skills_data[tag] + break + + assert skill_name + skill_value: int = self.get_skill_total_value(skill_name) + if weapon.quality == ItemQuality.EXCELLENT: + skill_value += 1 + + return replace( + weapon, + name=f"[{skill_value}] {weapon.name}") + + sorted_melee_weapon = sorted( + [add_skill_value(x) for x in self.weapons if "MeleeWeapon" in x.get_all_tags()], + key=weapon_sorter, + reverse=True) + + sorted_ranged_weapon = sorted( + [add_skill_value(x) for x in self.weapons if "RangedWeapon" in x.get_all_tags()], + key=weapon_sorter, + reverse=True) assert len(sorted_melee_weapon) + len(sorted_ranged_weapon) == len(self.weapons)