from wordcloud import WordCloud, STOPWORDS, ImageColorGenerator import matplotlib.pyplot as plt import wikipedia import sys import warnings # supressing unnecessary warnings warnings.filterwarnings("ignore") # function to search the wikipedia article and generate the wordcloud def gen_cloud(topic): try: content = str(wikipedia.page(topic).content) except: print("Error, try searching something else...") sys.exit() STOPWORDS.add('==') stopwords = set(STOPWORDS) wordcloud = WordCloud(stopwords=stopwords, max_words=200, background_color="black", width=600, height=350).generate(content) return wordcloud # function to save the wordcloud to current directory def save_cloud(wordcloud): wordcloud.to_file("./wordcloud.png") # function to display the wordcloud with matplotlib def show_cloud(wordcloud): plt.imshow(wordcloud, interpolation='bilinear') plt.axis("off") plt.show() # driver code if __name__ == '__main__': topic = input("What do you want to search: ").strip() wordcloud = gen_cloud(topic) save_cloud(wordcloud) print("Wordcloud saved to current directory as wordcloud.png") desc = input("Do you wish to see the output(y/n): ") if desc == 'y': show_cloud(wordcloud) sys.exit()