diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 7573778dfe..f539854827 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -2285,6 +2285,16 @@ dependencies = [ "ordermap", ] +[[package]] +name = "phf" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "928c6535de93548188ef63bb7c4036bd415cd8f36ad25af44b9789b2ee72a48c" +dependencies = [ + "phf_macros", + "phf_shared 0.11.1", +] + [[package]] name = "phf_generator" version = "0.7.24" @@ -2305,6 +2315,29 @@ dependencies = [ "rand 0.8.5", ] +[[package]] +name = "phf_generator" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1181c94580fa345f50f19d738aaa39c0ed30a600d95cb2d3e23f94266f14fbf" +dependencies = [ + "phf_shared 0.11.1", + "rand 0.8.5", +] + +[[package]] +name = "phf_macros" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92aacdc5f16768709a569e913f7451034034178b05bdc8acda226659a3dccc66" +dependencies = [ + "phf_generator 0.11.1", + "phf_shared 0.11.1", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "phf_shared" version = "0.7.24" @@ -2323,6 +2356,15 @@ dependencies = [ "siphasher 0.3.10", ] +[[package]] +name = "phf_shared" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1fb5f6f826b772a8d4c0394209441e7d37cbbb967ae9c7e0e8134365c9ee676" +dependencies = [ + "siphasher 0.3.10", +] + [[package]] name = "pin-project" version = "1.0.11" @@ -4563,6 +4605,7 @@ dependencies = [ "lettre", "magic-crypt", "mime_guess", + "phf", "prometheus", "rand 0.8.5", "rand_core 0.6.3", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index eb283e534a..3c15cc2d72 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -25,6 +25,7 @@ tracing = "^0" tracing-subscriber = { version = "^0", features = ["env-filter", "json"]} console-subscriber = "^0" prometheus = { version = "^0", default-features = false } +phf = { version = "0.11", features = ["macros"] } rust-embed = "^6" mime_guess = "^2" diff --git a/backend/src/parser.rs b/backend/src/parser.rs index db7a9e8159..369bb7c806 100644 --- a/backend/src/parser.rs +++ b/backend/src/parser.rs @@ -9,6 +9,7 @@ use std::collections::HashMap; use itertools::Itertools; +use phf::phf_map; use regex::Regex; use serde::Serialize; use serde_json::json; @@ -638,6 +639,18 @@ fn to_value(et: &ExpressionType) -> Option { } } +static PYTHON_IMPORTS_REPLACEMENT: phf::Map<&'static str, &'static str> = phf_map! { + "psycopg2" => "psycopg2-binary" +}; + +fn replace_import(x: &str) -> String { + PYTHON_IMPORTS_REPLACEMENT + .get(x) + .map(|x| x.to_owned()) + .unwrap_or(x) + .to_string() +} + pub fn parse_python_imports(code: &str) -> error::Result> { let find_requirements = code .lines() @@ -647,10 +660,8 @@ pub fn parse_python_imports(code: &str) -> error::Result> { let lines = code .lines() .skip(pos + 1) - .map_while(|x| { - re.captures(x) - .map(|x| x.get(1).unwrap().as_str().to_string()) - }) + .map_while(|x| re.captures(x).map(|x| x.get(1).unwrap().as_str())) + .map(replace_import) .collect(); Ok(lines) } else { @@ -670,12 +681,9 @@ pub fn parse_python_imports(code: &str) -> error::Result> { .collect::>(), ), StatementType::ImportFrom { level: _, module: Some(mod_), names: _ } => { - Some(vec![mod_ - .split('.') - .next() - .unwrap_or("") - .to_string() - .replace("_", "-")]) + let imprt = mod_.split('.').next().unwrap_or("").replace("_", "-"); + + Some(vec![replace_import(&imprt)]) } _ => None, },