53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
|
import random
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from random_sets.datasources import DataSource
|
||
|
|
|
||
|
|
|
||
|
|
class WeightedSet:
|
||
|
|
"""
|
||
|
|
A set in which members each have a weight, used for selecting at random.
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
>>> ws = WeightedSet(('foo', 1.0), ('bar', 0.5))
|
||
|
|
>>> ws.random()
|
||
|
|
('foo', 1.0)
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, *weighted_members: tuple):
|
||
|
|
self.members = []
|
||
|
|
self.weights = []
|
||
|
|
if weighted_members:
|
||
|
|
self.members, self.weights = list(zip(*weighted_members))
|
||
|
|
|
||
|
|
def random(self) -> str:
|
||
|
|
return random.choices(self.members, self.weights)[0]
|
||
|
|
|
||
|
|
def __add__(self, obj):
|
||
|
|
ws = WeightedSet()
|
||
|
|
ws.members = self.members + obj.members
|
||
|
|
ws.weights = self.weights + obj.weights
|
||
|
|
return ws
|
||
|
|
|
||
|
|
def __str__(self):
|
||
|
|
return f"{self.members}\n{self.weights}"
|
||
|
|
|
||
|
|
|
||
|
|
class DataSourceSet(WeightedSet):
|
||
|
|
|
||
|
|
def __init__(self, source: Path):
|
||
|
|
self.source = DataSource(source.read_text())
|
||
|
|
super().__init__(*[(key, value) for key, value in self.source.frequencies.items()])
|
||
|
|
|
||
|
|
def random(self):
|
||
|
|
random_key = super().random()
|
||
|
|
return self.source.as_dict()[random_key]
|
||
|
|
|
||
|
|
|
||
|
|
def equal_weights(terms: list, weight: float = 1.0, blank: bool = True) -> WeightedSet:
|
||
|
|
ws = WeightedSet(*[(term, weight) for term in terms])
|
||
|
|
if blank:
|
||
|
|
ws = WeightedSet(("", 1.0)) + ws
|
||
|
|
return ws
|