Merge pull request #133 from stijndcl/covid

Rewrite Covid

didier/cogs/other.py

@@ -7,9 +7,10 @@ from discord.ext import commands
 from database.crud.links import get_link_by_name
 from database.schemas import Link
 from didier import Didier
-from didier.data.apis import inspirobot, urban_dictionary
+from didier.data.apis import disease_sh, inspirobot, urban_dictionary
 from didier.data.embeds.google import GoogleSearch
 from didier.data.scrapers import google
+from didier.utils.discord.autocompletion.country import autocomplete_country
 
 
 class Other(commands.Cog):
@@ -20,16 +21,34 @@ class Other(commands.Cog):
     def __init__(self, client: Didier):
         self.client = client
 
+    @commands.hybrid_command(name="corona", aliases=["covid", "rona"])
+    async def covid(self, ctx: commands.Context, country: str = "Belgium"):
+        """Show Covid-19 info for a specific country.
+
+        By default, this will fetch the numbers for Belgium.
+
+        To get worldwide stats, use `all`, `global`, `world`, or `worldwide`.
+        """
+        async with ctx.typing():
+            if country.lower() in ["all", "global", "world", "worldwide"]:
+                data = await disease_sh.get_global_info(self.client.http_session)
+            else:
+                data = await disease_sh.get_country_info(self.client.http_session, country)
+
+            await ctx.reply(embed=data.to_embed(), mention_author=False)
+
+    @covid.autocomplete("country")
+    async def _covid_country_autocomplete(self, interaction: discord.Interaction, value: str):
+        """Autocompletion for the 'country'-parameter"""
+        return autocomplete_country(value)[:25]
+
     @commands.hybrid_command(
         name="define", aliases=["ud", "urban"], description="Look up the definition of a word on the Urban Dictionary"
     )
     async def define(self, ctx: commands.Context, *, query: str):
         """Look up the definition of `query` on the Urban Dictionary."""
         async with ctx.typing():
-            status_code, definitions = await urban_dictionary.lookup(self.client.http_session, query)
+            definitions = await urban_dictionary.lookup(self.client.http_session, query)
-            if not definitions:
-                return await ctx.reply(f"Something went wrong (status {status_code})")
-
             await ctx.reply(embed=definitions[0].to_embed(), mention_author=False)
 
     @commands.hybrid_command(name="google", description="Google search")

didier/data/apis/disease_sh.py

@@ -0,0 +1,38 @@
+from aiohttp import ClientSession
+
+from didier.data.embeds.disease_sh import CovidData
+from didier.utils.http.requests import ensure_get
+
+__all__ = ["get_country_info", "get_global_info"]
+
+
+async def get_country_info(http_session: ClientSession, country: str) -> CovidData:
+    """Fetch the info for a given country for today and yesterday"""
+    endpoint = f"https://disease.sh/v3/covid-19/countries/{country}"
+
+    params = {"yesterday": 0, "strict": 1, "allowNull": 0}
+    async with ensure_get(http_session, endpoint, params=params) as response:
+        today = response
+
+    params["yesterday"] = 1
+    async with ensure_get(http_session, endpoint, params=params) as response:
+        yesterday = response
+
+    data = {"today": today, "yesterday": yesterday}
+    return CovidData.parse_obj(data)
+
+
+async def get_global_info(http_session: ClientSession) -> CovidData:
+    """Fetch the global info for today and yesterday"""
+    endpoint = "https://disease.sh/v3/covid-19/all"
+
+    params = {"yesterday": 0, "allowNull": 0}
+    async with ensure_get(http_session, endpoint, params=params) as response:
+        today = response
+
+    params["yesterday"] = 1
+    async with ensure_get(http_session, endpoint, params=params) as response:
+        yesterday = response
+
+    data = {"today": today, "yesterday": yesterday}
+    return CovidData.parse_obj(data)

didier/data/apis/urban_dictionary.py

@@ -1,8 +1,7 @@
-from http import HTTPStatus
-
 from aiohttp import ClientSession
 
 from didier.data.embeds.urban_dictionary import Definition
+from didier.utils.http.requests import ensure_get
 
 __all__ = ["lookup", "PER_PAGE"]
 
@@ -10,13 +9,9 @@ __all__ = ["lookup", "PER_PAGE"]
 PER_PAGE = 10
 
 
-async def lookup(http_session: ClientSession, query: str) -> tuple[int, list[Definition]]:
+async def lookup(http_session: ClientSession, query: str) -> list[Definition]:
     """Fetch the Urban Dictionary definitions for a given word"""
     url = "https://api.urbandictionary.com/v0/define"
 
-    async with http_session.get(url, params={"term": query}) as response:
+    async with ensure_get(http_session, url, params={"term": query}) as response:
-        if response.status != HTTPStatus.OK:
+        return list(map(Definition.parse_obj, response["list"]))
-            return response.status, []
-
-        response_json = await response.json()
-        return 200, list(map(Definition.parse_obj, response_json["list"]))

didier/data/embeds/disease_sh.py

@@ -0,0 +1,100 @@
+import discord
+from overrides import overrides
+from pydantic import BaseModel, Field, validator
+
+from didier.data.embeds.base import EmbedPydantic
+
+__all__ = ["CovidData"]
+
+
+class _CovidNumbers(BaseModel):
+    """Covid numbers for a country
+
+    For worldwide numbers, country_info will be None
+    """
+
+    updated: int
+    country: str = "Worldwide"
+    cases: int
+    today_cases: int = Field(alias="todayCases")
+    deaths: int
+    today_deaths: int = Field(alias="todayDeaths")
+    recovered: int
+    today_recovered: int = Field(alias="todayRecovered")
+    active: int
+    tests: int
+
+    @validator("updated")
+    def updated_to_seconds(cls, value: int) -> int:
+        """Turn the updated field into seconds instead of milliseconds"""
+        return int(value) // 1000
+
+
+class CovidData(EmbedPydantic):
+    """Covid information from two days combined into one model"""
+
+    today: _CovidNumbers
+    yesterday: _CovidNumbers
+
+    @overrides
+    def to_embed(self, **kwargs) -> discord.Embed:
+        embed = discord.Embed(colour=discord.Colour.red(), title=f"Coronatracker {self.today.country}")
+        embed.description = f"Last update: <t:{self.today.updated}:R>"
+        embed.set_thumbnail(url="https://i.imgur.com/aWnDuBt.png")
+
+        cases_indicator = self._trend_indicator(self.today.today_cases, self.yesterday.today_cases)
+        embed.add_field(
+            name="Cases (Today)",
+            value=f"{self.today.cases:,} **({self.today.today_cases:,})** {cases_indicator}".replace(",", "."),
+            inline=False,
+        )
+
+        active_indicator = self._trend_indicator(self.today.active, self.yesterday.active)
+        active_diff = self.today.active - self.yesterday.active
+        embed.add_field(
+            name="Active (Today)",
+            value=f"{self.today.active:,} **({self._with_sign(active_diff)})** {active_indicator}".replace(",", "."),
+            inline=False,
+        )
+
+        deaths_indicator = self._trend_indicator(self.today.today_deaths, self.yesterday.today_deaths)
+        embed.add_field(
+            name="Deaths (Today)",
+            value=f"{self.today.deaths:,} **({self.today.today_deaths:,})** {deaths_indicator}".replace(",", "."),
+            inline=False,
+        )
+
+        recovered_indicator = self._trend_indicator(self.today.today_recovered, self.yesterday.today_recovered)
+        embed.add_field(
+            name="Recovered (Today)",
+            value=f"{self.today.recovered} **({self.today.today_recovered:,})** {recovered_indicator}".replace(
+                ",", "."
+            ),
+            inline=False,
+        )
+
+        tests_diff = self.today.tests - self.yesterday.tests
+        embed.add_field(
+            name="Tests Administered (Today)",
+            value=f"{self.today.tests:,} **({tests_diff:,})**".replace(",", "."),
+            inline=False,
+        )
+
+        return embed
+
+    def _with_sign(self, value: int) -> str:
+        """Prepend a + symbol if a number is positive"""
+        if value > 0:
+            return f"+{value:,}"
+
+        return f"{value:,}"
+
+    def _trend_indicator(self, today: int, yesterday: int) -> str:
+        """Function that returns a rise/decline indicator for the target key."""
+        if today > yesterday:
+            return ":small_red_triangle:"
+
+        if yesterday > today:
+            return ":small_red_triangle_down:"
+
+        return ""

didier/utils/discord/autocompletion/country.py

@@ -0,0 +1,250 @@
+from discord import app_commands
+
+__all__ = ["autocomplete_country"]
+
+# This list was parsed out of a request to https://disease.sh/v3/covid-19/countries
+country_list = [
+    "Afghanistan",
+    "Albania",
+    "Algeria",
+    "Andorra",
+    "Angola",
+    "Anguilla",
+    "Antigua and Barbuda",
+    "Argentina",
+    "Armenia",
+    "Aruba",
+    "Australia",
+    "Austria",
+    "Azerbaijan",
+    "Bahamas",
+    "Bahrain",
+    "Bangladesh",
+    "Barbados",
+    "Belarus",
+    "Belgium",
+    "Belize",
+    "Benin",
+    "Bermuda",
+    "Bhutan",
+    "Bolivia",
+    "Bosnia",
+    "Botswana",
+    "Brazil",
+    "British Virgin Islands",
+    "Brunei",
+    "Bulgaria",
+    "Burkina Faso",
+    "Burundi",
+    "Cabo Verde",
+    "Cambodia",
+    "Cameroon",
+    "Canada",
+    "Caribbean Netherlands",
+    "Cayman Islands",
+    "Central African Republic",
+    "Chad",
+    "Channel Islands",
+    "Chile",
+    "China",
+    "Colombia",
+    "Comoros",
+    "Congo",
+    "Cook Islands",
+    "Costa Rica",
+    "Croatia",
+    "Cuba",
+    "Curaçao",
+    "Cyprus",
+    "Czechia",
+    "Côte d'Ivoire",
+    "DRC",
+    "Denmark",
+    "Diamond Princess",
+    "Djibouti",
+    "Dominica",
+    "Dominican Republic",
+    "Ecuador",
+    "Egypt",
+    "El Salvador",
+    "Equatorial Guinea",
+    "Eritrea",
+    "Estonia",
+    "Ethiopia",
+    "Falkland Islands (Malvinas)",
+    "Faroe Islands",
+    "Fiji",
+    "Finland",
+    "France",
+    "French Guiana",
+    "French Polynesia",
+    "Gabon",
+    "Gambia",
+    "Georgia",
+    "Germany",
+    "Ghana",
+    "Gibraltar",
+    "Greece",
+    "Greenland",
+    "Grenada",
+    "Guadeloupe",
+    "Guatemala",
+    "Guinea",
+    "Guinea-Bissau",
+    "Guyana",
+    "Haiti",
+    "Holy See (Vatican City State)",
+    "Honduras",
+    "Hong Kong",
+    "Hungary",
+    "Iceland",
+    "India",
+    "Indonesia",
+    "Iran",
+    "Iraq",
+    "Ireland",
+    "Isle of Man",
+    "Israel",
+    "Italy",
+    "Jamaica",
+    "Japan",
+    "Jordan",
+    "Kazakhstan",
+    "Kenya",
+    "Kiribati",
+    "Kuwait",
+    "Kyrgyzstan",
+    "Lao People's Democratic Republic",
+    "Latvia",
+    "Lebanon",
+    "Lesotho",
+    "Liberia",
+    "Libyan Arab Jamahiriya",
+    "Liechtenstein",
+    "Lithuania",
+    "Luxembourg",
+    "MS Zaandam",
+    "Macao",
+    "Macedonia",
+    "Madagascar",
+    "Malawi",
+    "Malaysia",
+    "Maldives",
+    "Mali",
+    "Malta",
+    "Marshall Islands",
+    "Martinique",
+    "Mauritania",
+    "Mauritius",
+    "Mayotte",
+    "Mexico",
+    "Micronesia",
+    "Moldova",
+    "Monaco",
+    "Mongolia",
+    "Montenegro",
+    "Montserrat",
+    "Morocco",
+    "Mozambique",
+    "Myanmar",
+    "N. Korea",
+    "Namibia",
+    "Nauru",
+    "Nepal",
+    "Netherlands",
+    "New Caledonia",
+    "New Zealand",
+    "Nicaragua",
+    "Niger",
+    "Nigeria",
+    "Niue",
+    "Norway",
+    "Oman",
+    "Pakistan",
+    "Palau",
+    "Palestine",
+    "Panama",
+    "Papua New Guinea",
+    "Paraguay",
+    "Peru",
+    "Philippines",
+    "Poland",
+    "Portugal",
+    "Qatar",
+    "Romania",
+    "Russia",
+    "Rwanda",
+    "Réunion",
+    "S. Korea",
+    "Saint Helena",
+    "Saint Kitts and Nevis",
+    "Saint Lucia",
+    "Saint Martin",
+    "Saint Pierre Miquelon",
+    "Saint Vincent and the Grenadines",
+    "Samoa",
+    "San Marino",
+    "Sao Tome and Principe",
+    "Saudi Arabia",
+    "Senegal",
+    "Serbia",
+    "Seychelles",
+    "Sierra Leone",
+    "Singapore",
+    "Sint Maarten",
+    "Slovakia",
+    "Slovenia",
+    "Solomon Islands",
+    "Somalia",
+    "South Africa",
+    "South Sudan",
+    "Spain",
+    "Sri Lanka",
+    "St. Barth",
+    "Sudan",
+    "Suriname",
+    "Swaziland",
+    "Sweden",
+    "Switzerland",
+    "Syrian Arab Republic",
+    "Taiwan",
+    "Tajikistan",
+    "Tanzania",
+    "Thailand",
+    "Timor-Leste",
+    "Togo",
+    "Tonga",
+    "Trinidad and Tobago",
+    "Tunisia",
+    "Turkey",
+    "Turks and Caicos Islands",
+    "Tuvalu",
+    "UAE",
+    "UK",
+    "USA",
+    "Uganda",
+    "Ukraine",
+    "Uruguay",
+    "Uzbekistan",
+    "Vanuatu",
+    "Venezuela",
+    "Vietnam",
+    "Wallis and Futuna",
+    "Western Sahara",
+    "Yemen",
+    "Zambia",
+    "Zimbabwe",
+]
+
+
+def autocomplete_country(argument: str) -> list[app_commands.Choice]:
+    """Autocompletion for country names"""
+    argument = argument.lower()
+
+    global_autocomplete = ["Global"]
+
+    return [
+        app_commands.Choice(name=country, value=country)
+        for country in global_autocomplete + country_list
+        if argument in country.lower()
+    ]

didier/utils/http/requests.py

@@ -1,6 +1,6 @@
 import logging
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Optional
 
 from aiohttp import ClientResponse, ClientSession, ContentTypeError
 
@@ -19,10 +19,10 @@ def request_successful(response: ClientResponse) -> bool:
 
 @asynccontextmanager
 async def ensure_get(
-    http_session: ClientSession, endpoint: str, *, log_exceptions: bool = True
+    http_session: ClientSession, endpoint: str, *, params: Optional[dict] = None, log_exceptions: bool = True
 ) -> AsyncGenerator[dict, None]:
     """Context manager that automatically raises an exception if a GET-request fails"""
-    async with http_session.get(endpoint) as response:
+    async with http_session.get(endpoint, params=params) as response:
         try:
             content = await response.json()
         except ContentTypeError: