Hrochobot/hrochobot/cogs/news.py

from datetime import datetime
import discord
from discord.ext import commands
from discord.utils import get
from markdownify import markdownify
import re
from bs4 import BeautifulSoup

from hrochobot.utils.ksp_utils import ksp_feed, strip_id, KSP_URL
import hrochobot.utils.data as data

NEWS_JSON = "news"

async def get_news_ids():
    feed = await ksp_feed()
    return list(map(lambda e: e.id, feed.entries))

async def autocomplete_news_ids(ctx):
    value = ctx.value.lower()
    options = []
    for id_ in map(strip_id, await get_news_ids()):
        lid = id_.lower()
        if lid.startswith(value) or lid.split("_", 1)[1].startswith(value):
            options.append(id_)
    return options

def guess_color(title):
    """
    Automagically guess color of given post.
    Not always reliable as all things automagic.
    """
    def contains(*regexes):
        return any(re.search(regex, title) for regex in regexes)

    if contains(r"(\d+)-Z(\d+)", "začátečnic", "KSP-Z"):
        return discord.Color.green()
    elif contains(r"(\d+)-(\d+)", "seriál", "série", "KSP-H"):
        return discord.Color.blue()
    else:
        return discord.Color.dark_purple()

def format_entry(entry, author=None):
    content = "\n\n".join(map(lambda x: x.replace('\n', ' '), entry.summary.split("\n\n")))
    embed = discord.Embed(
        title=entry.title,
        url=entry.link,
        description=markdownify(content, strip=["img"]),
        color=guess_color(entry.title),
    )

    soup = BeautifulSoup(content, 'html.parser')
    img = soup.find('img')
    if img:
        embed.set_image(url=img['src'])

    if author:
        embed.set_author(name=author)

    embed.set_thumbnail(url=f"{KSP_URL}/img/hippo_head.png")

    date = datetime.fromisoformat(entry.published)
    embed.set_footer(text=date.strftime("%-d. %-m. %Y"))

    return embed

async def post_news(bot, guild, entry_id):
    news_json = data.load_guild_data(guild.id, NEWS_JSON)
    if "news_channel" not in news_json:
        return "News channel not set."

    channel = get(guild.channels, id=news_json["news_channel"])
    feed = await ksp_feed()
    entries_with_id = list(filter(lambda e: strip_id(e.id) == entry_id, feed.entries))
    if len(entries_with_id) == 0:
        return f"Entry with id ``{entry_id}`` not found."

    await channel.send(embed=format_entry(entries_with_id[0], author=feed.feed.author))
    return None


class News(commands.Cog):
    def __init__(self, bot):
        self.bot = bot
    news = discord.SlashCommandGroup(
        "news",
        "Commands for management of ksp news.",
        guild_only=True,
        checks=[commands.has_permissions(manage_guild=True)]
    )

    @news.command(description="Sets channel for posting news.")
    @discord.option("channel_id", str, description="Id of the channel for sending news.")
    async def set_channel(self, ctx, channel_id: str):
        try:
            channel_id = int(channel_id)
        except ValueError:
            return await ctx.respond(f"Channel id must be int.", ephemeral=True)

        if not (channel := get(ctx.guild.channels, id=channel_id)):
            return await ctx.respond(f"No channel with id ``{channel_id}``.", ephemeral=True)

        news_json = data.load_guild_data(ctx.guild.id, NEWS_JSON)
        news_json["news_channel"] = channel_id
        data.dump_guild_data(ctx.guild.id, NEWS_JSON, news_json)
        return await ctx.respond(f"News channel set to {channel.mention}.", ephemeral=True)

    @news.command(description="Posts news of given id to set channel.")
    @discord.option("id", str, description="Id of entry to send.", autocomplete=autocomplete_news_ids)
    async def post_news(self, ctx, id: int):
        await ctx.defer(ephemeral=True)
        err = await post_news(self.bot, ctx.guild, id)
        if err:
            return await ctx.respond(err, ephemeral=True)
        return await ctx.respond(f"News posted.")


def setup(bot):
    bot.add_cog(News(bot))