easy-ml/webui/manager.py

from typing import Dict

from gradio.components import Component


class Manager:
    def __init__(self) -> None:
        self._id_to_elem: Dict[str, "Component"] = {}
        self._elem_to_id: Dict["Component", str] = {}

    def add_elems(self, elem_dict: Dict[str, "Component"]) -> None:
        for elem_id, elem in elem_dict.items():
            self._id_to_elem[elem_id] = elem
            self._elem_to_id[elem] = elem_id

    def get_elem_by_id(self, elem_id: str) -> "Component":
        return self._id_to_elem[elem_id]

    def _get_id_by_elem(self, elem: "Component") -> str:
        return self._elem_to_id[elem]

    def get(self, data, key):
        return data[self.get_elem_by_id(key)]

    def get_all(self, data) -> Dict:
        all = {}
        for k, v in self._id_to_elem.items():
            all[k] = data[v]
        return all


engine = Manager()