diff options
Diffstat (limited to '')
| -rw-r--r-- | cogs/plot.py | 128 |
1 files changed, 128 insertions, 0 deletions
diff --git a/cogs/plot.py b/cogs/plot.py new file mode 100644 index 0000000..1516bc1 --- /dev/null +++ b/cogs/plot.py | |||
| @@ -0,0 +1,128 @@ | |||
| 1 | # This file is part of Senator Sperma | ||
| 2 | # Licensed under the GPL-3.0 License | ||
| 3 | # Copyright (C) 2026 Salmonik @ salmonik@tuta.io | ||
| 4 | # https://codeberg.org/salmonik/senator-sperma | ||
| 5 | |||
| 6 | import asyncio | ||
| 7 | import logging | ||
| 8 | import math | ||
| 9 | import os | ||
| 10 | import re | ||
| 11 | |||
| 12 | import discord | ||
| 13 | from discord.ext import commands | ||
| 14 | |||
| 15 | logger = logging.getLogger(__name__) | ||
| 16 | |||
| 17 | try: | ||
| 18 | import matplotlib | ||
| 19 | |||
| 20 | matplotlib.use("Agg") | ||
| 21 | import matplotlib.pyplot as plt | ||
| 22 | |||
| 23 | PLOT_AVAILABLE = True | ||
| 24 | except ImportError: | ||
| 25 | PLOT_AVAILABLE = False | ||
| 26 | logger.warning("matplotlib not available, plot command will be disabled") | ||
| 27 | |||
| 28 | |||
| 29 | def parse_expression(expr: str) -> str: | ||
| 30 | """Convert user expression (x^3) to Python (x**3)""" | ||
| 31 | expr = expr.replace("^", "**") | ||
| 32 | expr = re.sub(r"(\d)([x])", r"\1*\2", expr) | ||
| 33 | expr = re.sub(r"([x])(\d)", r"\1*\2", expr) | ||
| 34 | return expr | ||
| 35 | |||
| 36 | |||
| 37 | class Plot(commands.Cog): | ||
| 38 | def __init__(self, bot): | ||
| 39 | self.bot = bot | ||
| 40 | |||
| 41 | def _generate_plot(self, func_expr: str, output_path: str): | ||
| 42 | """Generate plot (runs in executor)""" | ||
| 43 | parsed_expr = parse_expression(func_expr) | ||
| 44 | |||
| 45 | x = [i * 0.1 for i in range(-100, 101)] | ||
| 46 | y = [] | ||
| 47 | |||
| 48 | for xi in x: | ||
| 49 | try: | ||
| 50 | result = eval( | ||
| 51 | parsed_expr, | ||
| 52 | { | ||
| 53 | "x": xi, | ||
| 54 | "math": math, | ||
| 55 | "sin": math.sin, | ||
| 56 | "cos": math.cos, | ||
| 57 | "tan": math.tan, | ||
| 58 | "sqrt": math.sqrt, | ||
| 59 | "log": math.log, | ||
| 60 | "exp": math.exp, | ||
| 61 | "pi": math.pi, | ||
| 62 | "e": math.e, | ||
| 63 | }, | ||
| 64 | ) | ||
| 65 | if math.isfinite(result): | ||
| 66 | y.append(result) | ||
| 67 | else: | ||
| 68 | y.append(float("nan")) | ||
| 69 | except Exception: | ||
| 70 | y.append(float("nan")) | ||
| 71 | |||
| 72 | fig, ax = plt.subplots(figsize=(10, 6)) | ||
| 73 | ax.plot(x, y, linewidth=2, color="#00ff00") | ||
| 74 | ax.set_title(f"y = {func_expr}", fontsize=14) | ||
| 75 | ax.set_xlabel("x", fontsize=12) | ||
| 76 | ax.set_ylabel("y", fontsize=12) | ||
| 77 | ax.grid(True, alpha=0.3) | ||
| 78 | ax.axhline(y=0, color="white", linewidth=0.5) | ||
| 79 | ax.axvline(x=0, color="white", linewidth=0.5) | ||
| 80 | ax.set_facecolor("#1a1a1a") | ||
| 81 | fig.patch.set_facecolor("#1a1a1a") | ||
| 82 | |||
| 83 | plt.savefig(output_path, dpi=100, bbox_inches="tight") | ||
| 84 | plt.close() | ||
| 85 | |||
| 86 | @commands.slash_command( | ||
| 87 | name="funkcja", description="Generuje wykres funkcji matematycznej" | ||
| 88 | ) | ||
| 89 | async def plot( | ||
| 90 | self, | ||
| 91 | ctx, | ||
| 92 | funkcja: str = discord.Option( | ||
| 93 | str, "Funkcja np. x^3, sin(x), x^2 + 2*x - 1", required=True | ||
| 94 | ), | ||
| 95 | ): | ||
| 96 | if not PLOT_AVAILABLE: | ||
| 97 | await ctx.respond( | ||
| 98 | "Generowanie wykresów nie jest dostępne (brak matplotlib)", | ||
| 99 | ephemeral=True, | ||
| 100 | ) | ||
| 101 | return | ||
| 102 | |||
| 103 | await ctx.defer() | ||
| 104 | |||
| 105 | temp_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "temp") | ||
| 106 | os.makedirs(temp_dir, exist_ok=True) | ||
| 107 | temp_path = os.path.join( | ||
| 108 | temp_dir, f"plot_{os.getpid()}_{asyncio.get_event_loop().time()}.png" | ||
| 109 | ) | ||
| 110 | |||
| 111 | try: | ||
| 112 | loop = asyncio.get_event_loop() | ||
| 113 | await loop.run_in_executor(None, self._generate_plot, funkcja, temp_path) | ||
| 114 | await ctx.respond(file=discord.File(temp_path, filename="plot.png")) | ||
| 115 | logger.info(f"Generated plot for function: {funkcja}") | ||
| 116 | except Exception as e: | ||
| 117 | logger.error(f"Plot generation error: {e}") | ||
| 118 | await ctx.respond(f"Błąd generowania wykresu: {str(e)}", ephemeral=True) | ||
| 119 | finally: | ||
| 120 | try: | ||
| 121 | if os.path.exists(temp_path): | ||
| 122 | os.unlink(temp_path) | ||
| 123 | except Exception as e: | ||
| 124 | logger.error(f"Error cleaning up temp file: {e}") | ||
| 125 | |||
| 126 | |||
| 127 | def setup(bot): | ||
| 128 | bot.add_cog(Plot(bot)) | ||
