diff --git a/random_sets/sets.py b/random_sets/sets.py index 0dc2cea..d17bcfa 100644 --- a/random_sets/sets.py +++ b/random_sets/sets.py @@ -38,8 +38,15 @@ class DataSourceSet(WeightedSet): def __init__(self, source: Path): self.source = DataSource(source.read_text()) + self._populate() + + def _populate(self): super().__init__(*[(key, value) for key, value in self.source.frequencies.items()]) + def set_frequency(self, frequency: str): + self.source.set_frequency(frequency) + self._populate() + def random(self): random_key = super().random() return self.source.as_dict()[random_key]