diff --git a/tests/test_generation.py b/tests/test_generation.py index 78301c4..9945c37 100644 --- a/tests/test_generation.py +++ b/tests/test_generation.py @@ -1,4 +1,4 @@ -from untitled import generation, rules +from untitled import PACKAGE_ROOT, generation, namegen, rules def test_cat_description_gen(): @@ -21,7 +21,12 @@ def test_cat_description_gen(): def test_generated_names_pass_filters(): + names = [] + with open(PACKAGE_ROOT / "assets" / "cat_names.txt") as f: + for i in f.readlines(): + names.append(i.strip()) + name_generator = namegen.NameGenerator(names, rules.validate_cat_name) for _ in range(10000): - name = generation.generate_name() + name = name_generator.generate_name() assert rules.validate_cat_name(name, auto_gen=True) is None assert any(c in "aeiouy" for c in name.lower()) diff --git a/untitled/generation.py b/untitled/generation.py index 2c763a6..5e652b8 100644 --- a/untitled/generation.py +++ b/untitled/generation.py @@ -1,7 +1,6 @@ import random -from collections import defaultdict -from untitled import PACKAGE_ROOT, content, rules +from untitled import content def generate_trait_sentence(traits): @@ -27,53 +26,3 @@ def generate_cat_choice(personality=None): def generate_cat_choices(n=5): personalities = random.sample(content.CAT_PERSONALITIES, n) return [generate_cat_choice(p) for p in personalities] - - -# NAME GEN MODEL - - -def _pad_names_for_model(names): - new_names = [] - for name in names: - new_names.append(f"<<{name}>") - return new_names - - -def _build_model(words): - model = defaultdict(list) - for word in words: - for pos in range(len(word) - 2): - model[word[pos] + word[pos + 1]].append(word[pos + 2]) - return model - - -def _make_name(model): - result = "<<" - while True: - window = result[-2:] - nxt = random.choice(model[window]) - if nxt == ">": - break - result += nxt - return result.removeprefix("<<") - - -def generate_name(): - while True: - name = _make_name(_model) - if ( - any(c in "aeiou" for c in name.lower()) - and rules.validate_cat_name(name, auto_gen=True) is None - ): - return name.capitalize() - - -def _load_raw_names(): - names = [] - with open(PACKAGE_ROOT / "assets" / "cat_names.txt") as f: - for name in f.readlines(): - names.append(name.strip()) - return names - - -_model = _build_model(_pad_names_for_model(_load_raw_names())) diff --git a/untitled/namegen.py b/untitled/namegen.py new file mode 100644 index 0000000..d75220a --- /dev/null +++ b/untitled/namegen.py @@ -0,0 +1,39 @@ +import random +from collections import defaultdict + +# NAME GEN MODEL + + +class NameGenerator: + def __init__(self, corpus, validate=None): + self.model = self.build_model(self.pad_names_for_model(corpus)) + self.validate = validate + + def pad_names_for_model(self, names): + new_names = [] + for name in names: + new_names.append(f"<<{name}>") + return new_names + + def build_model(self, words): + model = defaultdict(list) + for word in words: + for pos in range(len(word) - 2): + model[word[pos] + word[pos + 1]].append(word[pos + 2]) + return model + + def make_name(self, model): + result = "<<" + while True: + window = result[-2:] + nxt = random.choice(model[window]) + if nxt == ">": + break + result += nxt + return result.removeprefix("<<") + + def generate_name(self): + while True: + name = self.make_name(self.model) + if self.validate and self.validate(name, auto_gen=True) is None: + return name diff --git a/untitled/rules.py b/untitled/rules.py index ac2d984..5890aaf 100644 --- a/untitled/rules.py +++ b/untitled/rules.py @@ -11,6 +11,8 @@ def validate_cat_name(name, auto_gen=False): else: if len(name) < 4 or len(name) > 9: return "Your cat's name must be greater than 3 characters and below 9 characters long." + if not any(c in "aeiou" for c in name.lower()): + return "Your cat's name must contain a vowel." if not all(c in ALLOWED for c in name): return "Your cat's name can only have letters and numbers." if not any(c in string.ascii_letters for c in name): diff --git a/untitled/screens/adoption.py b/untitled/screens/adoption.py index 5d44d91..1777133 100644 --- a/untitled/screens/adoption.py +++ b/untitled/screens/adoption.py @@ -1,7 +1,12 @@ -from untitled import generation, model, rules, ui +from untitled import PACKAGE_ROOT, generation, model, namegen, rules, ui def adoption(): + names = [] + with open(PACKAGE_ROOT / "assets" / "cat_names.txt") as f: + for i in f.readlines(): + names.append(i.strip()) + name_generator = namegen.NameGenerator(names, rules.validate_cat_name) print("Welcome to the shelter!") while True: auto_name = "" @@ -25,7 +30,7 @@ def adoption(): break print(error) else: - auto_name = generation.generate_name() + auto_name = name_generator.generate_name().capitalize() if ui.confirm( f"Do you want to adopt {name}, {generation.generate_trait_sentence(choice).lower()}?"