import functools import random from typing import List, Optional import requests from gazpacho import Soup, get from rhyme_with_ai.utils import find_last_word def query_rhyme_words(sentence: str, n_rhymes: int, language:str="english") -> List[str]: """Returns a list of rhyme words for a sentence. Parameters ---------- sentence : Sentence that may end with punctuation n_rhymes : Maximum number of rhymes to return Returns ------- List[str] -- List of words that rhyme with the final word """ last_word = find_last_word(sentence) if language == "english": return query_datamuse_api(last_word, n_rhymes) elif language == "dutch": return mick_rijmwoordenboek(last_word, n_rhymes) else: raise NotImplementedError(f"Unsupported language ({language}) expected 'english' or 'dutch'.") def query_datamuse_api(word: str, n_rhymes: Optional[int] = None) -> List[str]: """Query the DataMuse API. Parameters ---------- word : Word to rhyme with n_rhymes : Max rhymes to return Returns ------- Rhyme words """ out = requests.get( "https://api.datamuse.com/words", params={"rel_rhy": word} ).json() words = [_["word"] for _ in out] if n_rhymes is None: return words return words[:n_rhymes] @functools.lru_cache(maxsize=128, typed=False) def mick_rijmwoordenboek(word: str, n_words: int): url = f"https://rijmwoordenboek.nl/rijm/{word}" html = get(url) soup = Soup(html) results = soup.find("div", {"id": "rhymeResultsWords"}).html.split("
") # clean up results = [r.replace("\n", "").replace(" ", "") for r in results] # filter html and empty strings results = [r for r in results if ("<" not in r) and (len(r) > 0)] return random.sample(results, min(len(results), n_words))