Skip to content

Commit efc0d8b

Browse files
authored
Add implicit returns to evals where last statement is an expr (#49)
1 parent 806a6fb commit efc0d8b

File tree

1 file changed

+20
-1
lines changed

1 file changed

+20
-1
lines changed

modules/eval.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,10 @@
2222
"""
2323
from __future__ import annotations
2424

25+
import ast
2526
import logging
2627
import textwrap
28+
import traceback
2729

2830
from discord.ext import commands
2931

@@ -56,7 +58,24 @@ def __init__(self, bot: core.Bot, endpoint_url: str) -> None:
5658
self.eval_endpoint: str = endpoint_url
5759

5860
async def perform_eval(self, code: core.Codeblock) -> str:
59-
formatted = CODE.format(user_code=textwrap.indent(code.content.replace("\t", " "), " "))
61+
source = code.content
62+
63+
try:
64+
parsed = ast.parse(source, filename="<input>")
65+
except SyntaxError as e:
66+
return "".join(traceback.format_exception(type(e), e, e.__traceback__))
67+
68+
if isinstance(parsed.body[-1], ast.Expr):
69+
expr: ast.Expr = parsed.body[-1]
70+
71+
lineno = expr.lineno - 1
72+
73+
co_mangled = source.splitlines()
74+
co_mangled[lineno] = "return " + co_mangled[lineno]
75+
76+
source = "\n".join(co_mangled)
77+
78+
formatted = CODE.format(user_code=textwrap.indent(source.replace("\t", " "), " "))
6079

6180
async with self.bot.session.post(self.eval_endpoint, json={"input": formatted}) as eval_response:
6281
if eval_response.status != 200:

0 commit comments

Comments
 (0)