from __future__ import annotations
from typing import TYPE_CHECKING, FrozenSet
import immutables
import yaml
if TYPE_CHECKING:
from nrel.hive.util.typealiases import EntityId, MembershipMap
[docs]def process_fleet_file(fleet_file: str, entity_type: str) -> MembershipMap:
"""
creates an immutable map that contains all of the fleet ids associated with the appropriate entity ids
:param fleet_file: the file to load memberships from
:param entity_type: the category the entity falls under, such as vehicles, bases, or stations
:return: an immutable map mapping all entity ids with the appropriate memberships
:raises Exception: from KEYErrors parsing the fleets file
"""
fleet_id_map: MembershipMap = immutables.Map()
fleet_id_map_mutation = fleet_id_map.mutate()
with open(fleet_file) as f:
config_dict = yaml.safe_load(f)
for fleet_id in config_dict:
for entity_id in config_dict[fleet_id][entity_type]:
if entity_id in fleet_id_map_mutation:
fleet_id_map_mutation[entity_id] = fleet_id_map_mutation[entity_id] + (
fleet_id,
)
else:
fleet_id_map_mutation.set(entity_id, (fleet_id,))
fleet_id_map = fleet_id_map_mutation.finish()
return fleet_id_map
[docs]def read_fleet_ids_from_file(fleet_file: str) -> FrozenSet[EntityId]:
"""
:param fleet_file: the file containing fleet information
:return: the list of base-level keys in the fleet file, which should be fleet ids
:raises: IOError if fleet_file does not exist or is not a yaml file
"""
try:
with open(fleet_file) as f:
fleet_dict = yaml.safe_load(f)
result = frozenset(fleet_dict.keys())
return result
except Exception as e:
raise IOError(f"failed reading fleet ids from fleet file {fleet_file}") from e