FeatureSet.py

class FeatureSet:

    def __init__(self, feature_set: dict, demographics:dict, statement_templates: dict, ):
        self.feature_set = feature_set
        self.feature_names = list(feature_set.keys())
        self.total_options = self.get_total_option_count()
        self.demographics = demographics
        self.statement_templates = statement_templates

    
    def get_feature_names(self) -> dict:
        return self.feature_names
    
    def get_feature_options(self, feature_name: str) -> list[str]:
        feature_dict = self.feature_set[feature_name]
        return list(feature_dict.keys())
    
    def get_feature_distribution(self, feature_name: str) -> dict:
        return self.feature_set[feature_name]
    
    def get_demographic_names(self) -> list[str]:
        return list(self.demographics.keys())
    
    def get_demographic_options(self, demographic: str) -> list[str]:
        demo_dict = self.demographics[demographic]
        return list(demo_dict.keys())
    
    def get_demographic_distribution(self, demographic: str) -> dict:
        return self.demographics[demographic]
        
    def get_total_option_count(self) -> int:
        sum = 0
        for feat in self.feature_names:
            sum += len(self.get_feature_options(feat))
        return sum
    
    def get_demographics(self) -> dict:
        return self.demographics

    def get_feature_fact(self, feature_name: str, feature_option: str) -> str:
            template = self.statement_templates[feature_name]
            return str(template.format(**{feature_name: feature_option}))
    
    def get_demographic_fact(self, demographic:str, demographic_option:str) -> str:
        template = self.statement_templates[demographic]
        return str(template.format(**{demographic: demographic_option}))