From 6dfc05217fe4ccb44b91371d0cfa3d9edf59567c Mon Sep 17 00:00:00 2001
From: classabbyamp <5366828+classabbyamp@users.noreply.github.com>
Date: Mon, 25 Jan 2021 17:24:45 -0500
Subject: [PATCH] add the ability to select an element to hamstudy (#344)

fixes #208
---
 CHANGELOG.md  |  1 +
 exts/study.py | 17 +++++++++++++++--
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3562965..dd32f06 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Added
 - MUF and foF2 maps from [prop.kc2g.com](https://prop.kc2g.com/).
 - Commands to show METAR (`?metar`) and TAF (`?taf`) (aeronautical weather conditions).
+- The ability to select an element of a pool in `?hamstudy`.
 ### Changed
 - New colour theme for `?greyline`.
 - Moved great circle distance and bearing calculation from `?ungrid` to `?griddistance`.
diff --git a/exts/study.py b/exts/study.py
index f55dd9a..d4fc5dc 100644
--- a/exts/study.py
+++ b/exts/study.py
@@ -31,13 +31,14 @@ class StudyCog(commands.Cog):
         self.session = aiohttp.ClientSession(connector=bot.qrm.connector)
 
     @commands.command(name="hamstudy", aliases=["rq", "randomquestion", "randomq"], category=cmn.cat.study)
-    async def _random_question(self, ctx: commands.Context, country: str = "", level: str = ""):
+    async def _random_question(self, ctx: commands.Context, country: str = "", level: str = "", element: str = ""):
         """Gets a random question from [HamStudy's](https://hamstudy.org) question pools."""
         with ctx.typing():
             embed = cmn.embed_factory(ctx)
 
             country = country.lower()
             level = level.lower()
+            element = element.upper()
 
             if country in study.pool_names.keys():
                 if level in study.pool_names[country].keys():
@@ -115,7 +116,19 @@ class StudyCog(commands.Cog):
                 pool = json.loads(await resp.read())["pool"]
 
             # Select a question
-            pool_section = random.choice(pool)["sections"]
+            if element:
+                els = [el["id"] for el in pool]
+                if element in els:
+                    pool_section = pool[els.index(element)]["sections"]
+                else:
+                    embed.title = "Element Not Found!"
+                    embed.description = f"Possible Elements for Country `{country}` and Level `{level}` are:"
+                    embed.colour = cmn.colours.bad
+                    embed.description += "\n\n" + "`" + "`, `".join(els) + "`"
+                    await ctx.send(embed=embed)
+                    return
+            else:
+                pool_section = random.choice(pool)["sections"]
             pool_questions = random.choice(pool_section)["questions"]
             question = random.choice(pool_questions)