base_generator.py

from abc import ABC, abstractmethod
from typing import ClassVar, List, Type, TypeVar, Generic

from rosnav_rl.states import SimulationStateContainer

from ..collectors import BaseUnit

GeneratedDataType = TypeVar("D")


class ObservationGeneratorUnit(BaseUnit, Generic[GeneratedDataType], ABC):
    """
    Base class for observation generator units.

    Attributes:
        name (str): The name of the observation generator unit.
        requires (List[BaseUnit]): A list of required base units.
        data_class (Type[D]): The type of data generated by the observation generator unit.
    """

    name: ClassVar[str]
    requires: ClassVar[List[BaseUnit]]
    data_class: Type[GeneratedDataType]

    @abstractmethod
    def generate(
        self,
        obs_dict: dict,
        simulation_state_container: SimulationStateContainer,
        *args,
        **kwargs,
    ) -> GeneratedDataType:
        """
        Generates the observation data based on the given observation dictionary.

        Args:
            obs_dict (dict): The observation dictionary.

        Returns:
            D: The generated observation data.
        """
        raise NotImplementedError

    # def generate_with_check(
    #     self,
    #     obs_dict: dict,
    #     simulation_state_container: SimulationStateContainer,
    #     *args,
    #     **kwargs,
    # ) -> GeneratedData:
    #     """
    #     Wrapper method that checks if 'obs_dict' contains elements from self.requires
    #     before calling the generate method.

    #     Args:
    #         obs_dict (dict): The observation dictionary.

    #     Returns:
    #         D: The generated observation data.
    #     """
    #     missing_requirements = [req for req in self.requires if req not in obs_dict]
    #     if missing_requirements:
    #         raise ValueError(f"Missing required observations: {missing_requirements}")

    #     return self.generate(obs_dict, simulation_state_container, *args, **kwargs)

    # def wrapper_generate_method(func):
    #     def wrapper(self, obs_dict, simulation_state_container, *args, **kwargs):
    #         missing_requirements = [req for req in self.requires if req not in obs_dict]
    #         if missing_requirements:
    #             raise ValueError(
    #                 f"Missing required observations: {missing_requirements}"
    #             )
    #         return func(self, obs_dict, simulation_state_container, *args, **kwargs)

    #     return wrapper

    # @wrapper_generate_method
    # def generate(self, obs_dict, simulation_state_container, *args, **kwargs):
    #     pass