diff options
Diffstat (limited to 'cogs/plot.py')
| -rw-r--r-- | cogs/plot.py | 65 |
1 files changed, 32 insertions, 33 deletions
diff --git a/cogs/plot.py b/cogs/plot.py index 1516bc1..65dcae4 100644 --- a/cogs/plot.py +++ b/cogs/plot.py | |||
| @@ -14,20 +14,28 @@ from discord.ext import commands | |||
| 14 | 14 | ||
| 15 | logger = logging.getLogger(__name__) | 15 | logger = logging.getLogger(__name__) |
| 16 | 16 | ||
| 17 | try: | 17 | _plt_module = None |
| 18 | import matplotlib | 18 | _plot_available = False |
| 19 | 19 | ||
| 20 | matplotlib.use("Agg") | 20 | |
| 21 | import matplotlib.pyplot as plt | 21 | def _ensure_plot(): |
| 22 | 22 | global _plt_module, _plot_available | |
| 23 | PLOT_AVAILABLE = True | 23 | |
| 24 | except ImportError: | 24 | if _plt_module is not None: |
| 25 | PLOT_AVAILABLE = False | 25 | return _plot_available |
| 26 | logger.warning("matplotlib not available, plot command will be disabled") | 26 | |
| 27 | try: | ||
| 28 | import matplotlib | ||
| 29 | matplotlib.use("Agg") | ||
| 30 | import matplotlib.pyplot as _plt_module | ||
| 31 | _plot_available = True | ||
| 32 | except ImportError: | ||
| 33 | _plot_available = False | ||
| 34 | logger.warning("matplotlib not available, plot command will be disabled") | ||
| 35 | return _plot_available | ||
| 27 | 36 | ||
| 28 | 37 | ||
| 29 | def parse_expression(expr: str) -> str: | 38 | def parse_expression(expr: str) -> str: |
| 30 | """Convert user expression (x^3) to Python (x**3)""" | ||
| 31 | expr = expr.replace("^", "**") | 39 | expr = expr.replace("^", "**") |
| 32 | expr = re.sub(r"(\d)([x])", r"\1*\2", expr) | 40 | expr = re.sub(r"(\d)([x])", r"\1*\2", expr) |
| 33 | expr = re.sub(r"([x])(\d)", r"\1*\2", expr) | 41 | expr = re.sub(r"([x])(\d)", r"\1*\2", expr) |
| @@ -39,7 +47,12 @@ class Plot(commands.Cog): | |||
| 39 | self.bot = bot | 47 | self.bot = bot |
| 40 | 48 | ||
| 41 | def _generate_plot(self, func_expr: str, output_path: str): | 49 | def _generate_plot(self, func_expr: str, output_path: str): |
| 42 | """Generate plot (runs in executor)""" | 50 | global _plt_module |
| 51 | |||
| 52 | if not _ensure_plot(): | ||
| 53 | raise RuntimeError("matplotlib not available") | ||
| 54 | |||
| 55 | plt = _plt_module | ||
| 43 | parsed_expr = parse_expression(func_expr) | 56 | parsed_expr = parse_expression(func_expr) |
| 44 | 57 | ||
| 45 | x = [i * 0.1 for i in range(-100, 101)] | 58 | x = [i * 0.1 for i in range(-100, 101)] |
| @@ -83,30 +96,16 @@ class Plot(commands.Cog): | |||
| 83 | plt.savefig(output_path, dpi=100, bbox_inches="tight") | 96 | plt.savefig(output_path, dpi=100, bbox_inches="tight") |
| 84 | plt.close() | 97 | plt.close() |
| 85 | 98 | ||
| 86 | @commands.slash_command( | 99 | @commands.slash_command(name="funkcja", description="Generuje wykres funkcji matematycznej") |
| 87 | name="funkcja", description="Generuje wykres funkcji matematycznej" | 100 | async def plot(self, ctx, funkcja: str = discord.Option(str, "Funkcja np. x^3, sin(x), x^2 + 2*x - 1", required=True)): |
| 88 | ) | 101 | if not _ensure_plot(): |
| 89 | async def plot( | 102 | await ctx.respond("Generowanie wykresów nie jest dostępne (brak matplotlib)", ephemeral=True) |
| 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 | 103 | return |
| 102 | 104 | ||
| 103 | await ctx.defer() | 105 | await ctx.defer() |
| 104 | 106 | ||
| 105 | temp_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "temp") | 107 | from cogs.utils.temp import get_temp_dir |
| 106 | os.makedirs(temp_dir, exist_ok=True) | 108 | temp_path = os.path.join(get_temp_dir(), f"plot_{os.getpid()}_{asyncio.get_event_loop().time()}.png") |
| 107 | temp_path = os.path.join( | ||
| 108 | temp_dir, f"plot_{os.getpid()}_{asyncio.get_event_loop().time()}.png" | ||
| 109 | ) | ||
| 110 | 109 | ||
| 111 | try: | 110 | try: |
| 112 | loop = asyncio.get_event_loop() | 111 | loop = asyncio.get_event_loop() |
| @@ -125,4 +124,4 @@ class Plot(commands.Cog): | |||
| 125 | 124 | ||
| 126 | 125 | ||
| 127 | def setup(bot): | 126 | def setup(bot): |
| 128 | bot.add_cog(Plot(bot)) | 127 | bot.add_cog(Plot(bot)) \ No newline at end of file |
