Allow images to be embedded from document

2 files changed, 16 insertions, 3 deletions

diff --git a/src/main.py b/src/main.py
index 88d2a4d..51fc8c5 100644
--- a/src/main.py
+++ b/src/main.py
@@ -25,7 +25,7 @@ async def explain(ctx, *args):
 
     try:
         search_content = scrape_lexicanum.get_search_response(config, query)
-        page_header, page_content = scrape_lexicanum.get_page_content(config, search_content[0])
+        page_header, page_content, page_img_link = scrape_lexicanum.get_page_content(config, search_content[0])
         
         embed = discord.Embed(
                     title=page_header,
@@ -38,6 +38,9 @@ async def explain(ctx, *args):
         embed.add_field(name="Overview", value=page_content)
         embed.add_field(name="You May Have Meant", value=string_results)
 
+        if page_img_link:
+            embed.set_image(url=page_img_link)
+
         await ctx.send(embed=embed)
     except scrape_lexicanum.WikiError as e:
         await ctx.send(f"{e}")
diff --git a/src/scrape_lexicanum.py b/src/scrape_lexicanum.py
index a719fb1..88001a3 100644
--- a/src/scrape_lexicanum.py
+++ b/src/scrape_lexicanum.py
@@ -7,6 +7,9 @@ import requests
 class WikiError(Exception):
     pass
 
+def fix_link(config, node, field):
+    node[field] = node[field].replace("/wiki", config["site"] + "/wiki")
+
 def get_search_response(config, query):
     try:
         response = requests.get(config['site'] + '/mediawiki/api.php?action=opensearch&format=json&formatversion=2&search=' + query + '&namespace=0&limit=6&suggest=true', timeout=7)
@@ -49,10 +52,17 @@ def get_page_content(config, page_name):
     parser_output = soup.find(id="mw-content-text").contents[0]
     content = parser_output.find("p")
     for a in content.findAll("a"):
-        a["href"] = a["href"].replace("/wiki", config["site"] + "/wiki")
+        fix_link(config, a, "href")
+
+    image = parser_output.find("a", {"class": "image"})
+    image_link = None
+    if image:
+        if image.contents[0]:
+            fix_link(config, image.contents[0], "src")
+            image_link = image.contents[0]["src"]
 
     header_md = markdownify(str(soup.find(id="firstHeading").contents[0]))
     content_md = markdownify(str(content))
 
-    return header_md, content_md
+    return header_md, content_md, image_link