kalugin66/windmill
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

optimize python parsing + adding tests

Browse Source
This commit is contained in:
Ruben Fiszel
2022-08-21 10:50:59 +02:00
parent c9e0afd652
commit f1c8f593da
3 changed files with 236 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
backend/src/parser.rs
View File

@@ -8,21 +8,21 @@
use serde::Serialize;
#[derive(Serialize)]
#[derive(Serialize, Debug, PartialEq)]
pub struct MainArgSignature {
pub star_args: bool,
pub star_kwargs: bool,
pub args: Vec<Arg>,
}
#[derive(Serialize, Clone)]
#[derive(Serialize, Clone, Debug, PartialEq)]
#[serde(rename_all(serialize = "lowercase"))]
pub struct ObjectProperty {
pub key: String,
pub typ: Box<Typ>,
}
#[derive(Serialize, Clone)]
#[derive(Serialize, Clone, Debug, PartialEq)]
#[serde(rename_all(serialize = "lowercase"))]
pub enum Typ {
Str(Option<Vec<String>>),
@@ -39,7 +39,7 @@ pub enum Typ {
Unknown,
}
#[derive(Serialize, Clone)]
#[derive(Serialize, Clone, Debug, PartialEq)]
pub struct Arg {
pub name: String,
pub typ: Typ,

171
backend/src/parser_py.rs
View File

@@ -23,8 +23,50 @@ use rustpython_parser::{
parser,
};
fn filter_non_main(code: &str) -> String {
const DEF_MAIN: &str = "def main(";
let mut filtered_code = String::new();
let mut code_iter = code.split("\n");
let mut remaining: String = String::new();
while let Some(line) = code_iter.next() {
if line.starts_with(DEF_MAIN) {
filtered_code += DEF_MAIN;
remaining += line.strip_prefix(DEF_MAIN).unwrap();
remaining += &code_iter.join("\n");
break;
}
}
if filtered_code.is_empty() {
return String::new();
}
let mut chars = remaining.chars();
let mut open_parens = 1;
while let Some(c) = chars.next() {
if c == '(' {
open_parens += 1;
} else if c == ')' {
open_parens -= 1;
}
filtered_code.push(c);
if open_parens == 0 {
break;
}
}
filtered_code.push_str(": return");
return filtered_code;
}
pub fn parse_python_signature(code: &str) -> error::Result<MainArgSignature> {
let ast = parser::parse_program(code)
let filtered_code = filter_non_main(code);
if filtered_code.is_empty() {
return Err(error::Error::BadRequest(
"No main function found".to_string(),
));
}
let ast = parser::parse_program(&filtered_code)
.map_err(|e| error::Error::ExecutionErr(format!("Error parsing code: {}", e.to_string())))?
.statements;
let param = ast.into_iter().find_map(|x| match x {
@@ -163,11 +205,16 @@ pub fn parse_python_imports(code: &str) -> error::Result<Vec<String>> {
.collect();
Ok(lines)
} else {
let code = &code
.split("\n")
.filter(|x| x.starts_with("import ") || x.starts_with("from "))
.join("\n");
let ast = parser::parse_program(code)
.map_err(|e| {
error::Error::ExecutionErr(format!("Error parsing code: {}", e.to_string()))
})?
.statements;
let imports = ast
.into_iter()
.filter_map(|x| match x {
@@ -191,6 +238,7 @@ pub fn parse_python_imports(code: &str) -> error::Result<Vec<String>> {
.filter(|x| !STDIMPORTS.contains(&x.as_str()))
.unique()
.collect();
Ok(imports)
}
}
@@ -198,12 +246,10 @@ pub fn parse_python_imports(code: &str) -> error::Result<Vec<String>> {
#[cfg(test)]
mod tests {
// Note this useful idiom: importing names from outer (for mod tests) scope.
use super::*;
#[test]
fn test_parse_python_sig() -> anyhow::Result<()> {
//let code = "print(2 + 3, fd=sys.stderr)";
let code = "
import os
@@ -215,7 +261,124 @@ def main(test1: str, name: datetime.datetime = datetime.now(), byte: bytes = byt
return {\"len\": len(name), \"splitted\": name.split() }
";
println!("{}", serde_json::to_string(&parse_python_signature(code)?)?);
//println!("{}", serde_json::to_string()?);
assert_eq!(
parse_python_signature(code)?,
MainArgSignature {
star_args: false,
star_kwargs: false,
args: vec![
Arg {
name: "test1".to_string(),
typ: Typ::Str(None),
default: None,
has_default: false
},
Arg {
name: "name".to_string(),
typ: Typ::Unknown,
default: Some(json!("<function call>")),
has_default: true
},
Arg {
name: "byte".to_string(),
typ: Typ::Bytes,
default: Some(json!("<function call>")),
has_default: true
}
]
}
);
Ok(())
}
#[test]
fn test_parse_python_sig_2() -> anyhow::Result<()> {
let code = "
import os
def main(test1: str,
name: datetime.datetime = datetime.now(),
byte: bytes = bytes(1)):
print(f\"Hello World and a warm welcome especially to {name}\")
print(\"The env variable at `all/pretty_secret`: \", os.environ.get(\"ALL_PRETTY_SECRET\"))
return {\"len\": len(name), \"splitted\": name.split() }
";
//println!("{}", serde_json::to_string()?);
assert_eq!(
parse_python_signature(code)?,
MainArgSignature {
star_args: false,
star_kwargs: false,
args: vec![
Arg {
name: "test1".to_string(),
typ: Typ::Str(None),
default: None,
has_default: false
},
Arg {
name: "name".to_string(),
typ: Typ::Unknown,
default: Some(json!("<function call>")),
has_default: true
},
Arg {
name: "byte".to_string(),
typ: Typ::Bytes,
default: Some(json!("<function call>")),
has_default: true
}
]
}
);
Ok(())
}
#[test]
fn test_parse_python_sig_3() -> anyhow::Result<()> {
let code = "
import os
def main(test1: str,
name: datetime.datetime = datetime.now(),
byte: bytes = bytes(1)): return
";
//println!("{}", serde_json::to_string()?);
assert_eq!(
parse_python_signature(code)?,
MainArgSignature {
star_args: false,
star_kwargs: false,
args: vec![
Arg {
name: "test1".to_string(),
typ: Typ::Str(None),
default: None,
has_default: false
},
Arg {
name: "name".to_string(),
typ: Typ::Unknown,
default: Some(json!("<function call>")),
has_default: true
},
Arg {
name: "byte".to_string(),
typ: Typ::Bytes,
default: Some(json!("<function call>")),
has_default: true
}
]
}
);
Ok(())
}

65
backend/src/worker.rs
View File

@@ -576,6 +576,7 @@ async fn handle_nondep_job(
workspace_id = %job.workspace_id,
"started setup python dependencies"
);
let child = if !disable_nsjail {
Command::new(nsjail_path)
.current_dir(job_dir)
@@ -745,6 +746,7 @@ print(res_json)
.stderr(Stdio::piped())
.spawn()?
};
*status = handle_child(job, db, logs, last_line, timeout, child).await;
tracing::info!(
worker_name = %worker_name,
@@ -1379,6 +1381,69 @@ mod tests {
}
}
#[sqlx::test(fixtures("base"))]
async fn test_python_flow_2(db: DB) {
initialize_tracing().await;
let flow: FlowValue = serde_json::from_value(serde_json::json!({
"modules": [
{
"value": {
"type": "rawscript",
"content": "import os\nimport wmill\nfrom datetime import datetime\n\nimport re\nimport requests\nfrom datetime import datetime\nfrom html import unescape\n\n\ndef extract(html):\n r = r\"<img class=\\\"[^\\\"]+ img-comic\\\"(([^>]*alt=\\\"(?P<alt>[^\\\"]+)\\\")|([^>]*src=\\\"(?P<url>[^\\\"]+)\\\"))+[^>]*>\"\n match = re.search(r, html)\n if match:\n return {'url': match.group(\"url\"), 'desc': unescape(match.group(\"alt\"))}\n return None\n\n\ndef fetch(date=None):\n if date is None:\n date = datetime.now()\n url = f\"https://dilbert.com/strip/{date:%Y-%m-%d}\"\n return requests.get(url, allow_redirects=False).text\n\n\ndef get_today_comic(date=None):\n return extract(fetch(date))\n\n\n# Our webeditor includes a syntax, type checker through a language server running pyright\n# and the autoformatter Black in our servers. Use Cmd/Ctrl + S to autoformat the code.\n# Beware that the code is only saved when you click Save and not across reload.\n# You can however navigate to any steps safely.\n\"\"\"\nThe client is used to interact with windmill itself through its standard API.\nOne can explore the methods available through autocompletion of `client.XXX`.\nOnly the most common methods are included for ease of use. Request more as\nfeedback if you feel you are missing important ones.\n\"\"\"\n\n\ndef main(\n date: str = None\n):\n dateToFetch = datetime.strptime(date, '%Y-%m-%d') if (date is not None and len(date)>0) else datetime.now()\n print(f\"Fetchind Dilber from {dateToFetch}\")\n # retrieve variables, including secrets by querying the windmill platform.\n # secret fetching is audited by windmill.\n # secret = wmill.get_variable(\"g/all/pretty_secret\")\n return get_today_comic(dateToFetch)\n",
"language": "python3"
},
"input_transform": {
"date": {
"expr": "undefined",
"type": "javascript"
}
}
},
{
"value": {
"type": "rawscript",
"content": "\nimport requests\n\n\ndef main(bot_token: str, chat_id: str, url: str, caption: str = \"\"):\n return requests.get(\n f\"https://api.telegram.org/bot{bot_token}/sendPhoto\",\n params={\"chat_id\": chat_id, \"photo\": url, \"caption\": caption},\n ).url\n",
"language": "python3"
},
"input_transform": {
"url": {
"type": "static",
"value": null
},
"caption": {
"type": "static",
"value": ""
},
"chat_id": {
"type": "static",
"value": null
},
"bot_token": {
"type": "static",
"value": null
}
}
}
]
})).unwrap();
for i in 0..10 {
println!("python flow iteration: {}", i);
let result = run_job_in_new_worker_until_complete(
&db,
JobPayload::RawFlow { value: flow.clone(), path: None },
)
.await;
assert_eq!(
result,
serde_json::json!("https://api.telegram.org/botNone/sendPhoto?caption="),
"iteration: {i}"
);
}
}
#[sqlx::test(fixtures("base"))]
async fn test_python_job(db: DB) {
initialize_tracing().await;