feat: added new validation functions and tests

This commit is contained in:
Gabriel Almeida 2023-03-28 20:41:18 -03:00
commit 8890d37a44
3 changed files with 183 additions and 8 deletions

View file

@ -24,7 +24,9 @@ class Function(BaseModel):
def get_function(self):
"""Get the function"""
return validate.eval_function(self.code)
function_name = validate.extract_function_name(self.code)
return validate.create_function(self.code, function_name)
class PythonFunction(Function):

View file

@ -4,6 +4,15 @@ import types
from typing import Dict
def add_type_ignores():
if not hasattr(ast, "TypeIgnore"):
class TypeIgnore(ast.AST):
_fields = ()
ast.TypeIgnore = TypeIgnore
def validate_code(code):
# Initialize the errors dictionary
errors = {"imports": {"errors": []}, "function": {"errors": []}}
@ -16,12 +25,7 @@ def validate_code(code):
return errors
# Add a dummy type_ignores field to the AST
if not hasattr(ast, "TypeIgnore"):
class TypeIgnore(ast.AST):
_fields = ()
ast.TypeIgnore = TypeIgnore
add_type_ignores()
tree.type_ignores = []
# Evaluate the import statements
@ -61,3 +65,105 @@ def eval_function(function_string: str):
if function_object is None:
raise ValueError("Function string does not contain a function")
return function_object
def execute_function(code, function_name, *args, **kwargs):
add_type_ignores()
module = ast.parse(code)
exec_globals = globals().copy()
for node in module.body:
if isinstance(node, ast.Import):
for alias in node.names:
try:
exec(
f"{alias.asname or alias.name} = importlib.import_module('{alias.name}')",
exec_globals,
locals(),
)
exec_globals[alias.asname or alias.name] = importlib.import_module(
alias.name
)
except ModuleNotFoundError as e:
raise ModuleNotFoundError(
f"Module {alias.name} not found. Please install it and try again."
) from e
function_code = next(
node
for node in module.body
if isinstance(node, ast.FunctionDef) and node.name == function_name
)
function_code.parent = None
code_obj = compile(
ast.Module(body=[function_code], type_ignores=[]), "<string>", "exec"
)
try:
exec(code_obj, exec_globals, locals())
except Exception as e:
# handle execution error here
pass
# Add the function to the exec_globals dictionary
exec_globals[function_name] = locals()[function_name]
return exec_globals[function_name](*args, **kwargs)
def create_function(code, function_name):
if not hasattr(ast, "TypeIgnore"):
class TypeIgnore(ast.AST):
_fields = ()
ast.TypeIgnore = TypeIgnore
module = ast.parse(code)
exec_globals = globals().copy()
for node in module.body:
if isinstance(node, ast.Import):
for alias in node.names:
try:
exec_globals[alias.asname or alias.name] = importlib.import_module(
alias.name
)
except ModuleNotFoundError as e:
raise ModuleNotFoundError(
f"Module {alias.name} not found. Please install it and try again."
) from e
function_code = next(
node
for node in module.body
if isinstance(node, ast.FunctionDef) and node.name == function_name
)
function_code.parent = None
code_obj = compile(
ast.Module(body=[function_code], type_ignores=[]), "<string>", "exec"
)
try:
exec(code_obj, exec_globals, locals())
except Exception as e:
pass
exec_globals[function_name] = locals()[function_name]
# Return a function that imports necessary modules and calls the target function
def wrapped_function(*args, **kwargs):
for module_name, module in exec_globals.items():
if isinstance(module, type(importlib)):
globals()[module_name] = module
return exec_globals[function_name](*args, **kwargs)
return wrapped_function
def extract_function_name(code):
module = ast.parse(code)
for node in module.body:
if isinstance(node, ast.FunctionDef):
return node.name
raise ValueError("No function definition found in the code string")