From edd193a1190f64c719abcd07f957a5360d10464f Mon Sep 17 00:00:00 2001
From: GreasySlug <9619abgoni@gmail.com>
Date: Sun, 4 Feb 2024 23:51:45 +0900
Subject: [PATCH] refactor: wrap and handle the py ver retrieval

---
 build.rs                         |  4 +++-
 crates/erg_common/config.rs      |  2 +-
 crates/erg_common/python_util.rs | 12 ++++++------
 crates/erg_compiler/codegen.rs   |  7 ++++++-
 tests/embed.rs                   |  4 ++--
 tests/test.rs                    |  4 ++--
 6 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/build.rs b/build.rs
index 96aa2273a..ada05ea83 100644
--- a/build.rs
+++ b/build.rs
@@ -1,7 +1,9 @@
 use erg_common::python_util::{env_magic_number, env_python_version};
 
 fn main() -> std::io::Result<()> {
-    let version = env_python_version();
+    let Some(version) = env_python_version() else {
+        panic!("Failed to get python version");
+    };
     if version.major != 3 {
         panic!("Python 3 is required");
     }
diff --git a/crates/erg_common/config.rs b/crates/erg_common/config.rs
index 1407d2388..254202a81 100644
--- a/crates/erg_common/config.rs
+++ b/crates/erg_common/config.rs
@@ -376,7 +376,7 @@ impl ErgConfig {
                         .parse::<String>()
                         .expect("the value of `-py-command` is not a valid Python command");
                     cfg.py_magic_num = Some(detect_magic_number(&py_command));
-                    cfg.target_version = Some(get_python_version(&py_command));
+                    cfg.target_version = get_python_version(&py_command);
                     cfg.py_command = Some(Box::leak(py_command.into_boxed_str()));
                 }
                 "--hex-py-magic-num" | "--hex-python-magic-number" => {
diff --git a/crates/erg_common/python_util.rs b/crates/erg_common/python_util.rs
index e3cd9ba3f..a6dbf91dd 100644
--- a/crates/erg_common/python_util.rs
+++ b/crates/erg_common/python_util.rs
@@ -693,7 +693,7 @@ impl std::str::FromStr for PythonVersion {
     }
 }
 
-pub fn get_python_version(py_command: &str) -> PythonVersion {
+pub fn get_python_version(py_command: &str) -> Option<PythonVersion> {
     let out = if cfg!(windows) {
         Command::new("cmd")
             .arg("/C")
@@ -710,19 +710,19 @@ pub fn get_python_version(py_command: &str) -> PythonVersion {
             .expect("cannot get the python version")
     };
     let s_version = String::from_utf8(out.stdout).unwrap();
-    let mut iter = s_version.split(' ');
-    let mut iter = iter.nth(1).unwrap().split('.');
+    let iter = s_version.split(' ').nth(1)?;
+    let mut iter = iter.split('.');
     let major = iter.next().and_then(|i| i.parse().ok()).unwrap_or(3);
     let minor = iter.next().and_then(|i| i.parse().ok());
     let micro = iter.next().and_then(|i| i.trim_end().parse().ok());
-    PythonVersion {
+    Some(PythonVersion {
         major,
         minor,
         micro,
-    }
+    })
 }
 
-pub fn env_python_version() -> PythonVersion {
+pub fn env_python_version() -> Option<PythonVersion> {
     get_python_version(&which_python())
 }
 
diff --git a/crates/erg_compiler/codegen.rs b/crates/erg_compiler/codegen.rs
index dcda23fd5..55859fa50 100644
--- a/crates/erg_compiler/codegen.rs
+++ b/crates/erg_compiler/codegen.rs
@@ -228,7 +228,12 @@ pub struct PyCodeGenerator {
 impl PyCodeGenerator {
     pub fn new(cfg: ErgConfig) -> Self {
         Self {
-            py_version: cfg.target_version.unwrap_or_else(env_python_version),
+            py_version: cfg.target_version.unwrap_or_else(|| {
+                let Some(version) = env_python_version() else {
+                    panic!("Failed to get python version");
+                };
+                version
+            }),
             cfg,
             str_cache: CacheSet::new(),
             prelude_loaded: false,
diff --git a/tests/embed.rs b/tests/embed.rs
index 9a500ec94..ff870881d 100644
--- a/tests/embed.rs
+++ b/tests/embed.rs
@@ -69,7 +69,7 @@ while! do! c < 7, do!:
 
 #[test]
 fn test_transpiler_embedding3() -> Result<(), ()> {
-    if env_python_version().minor < Some(10) {
+    if env_python_version().unwrap().minor < Some(10) {
         println!("skipped: {}", fn_name!());
         return Ok(());
     }
@@ -96,7 +96,7 @@ print!(i, end:=\"\")
 
 #[test]
 fn test_transpiler_embedding4() -> Result<(), ()> {
-    if env_python_version().minor < Some(10) {
+    if env_python_version().unwrap().minor < Some(10) {
         println!("skipped: {}", fn_name!());
         return Ok(());
     }
diff --git a/tests/test.rs b/tests/test.rs
index 9cd4b9bac..da562e8e8 100644
--- a/tests/test.rs
+++ b/tests/test.rs
@@ -172,7 +172,7 @@ fn exec_fib() -> Result<(), ()> {
 #[test]
 fn exec_helloworld() -> Result<(), ()> {
     // HACK: When running the test with Windows, the exit code is 1 (the cause is unknown)
-    if cfg!(windows) && env_python_version().minor >= Some(8) {
+    if cfg!(windows) && env_python_version().unwrap().minor >= Some(8) {
         expect_end_with("examples/helloworld.er", 0, 1)
     } else {
         expect_success("examples/helloworld.er", 0)
@@ -323,7 +323,7 @@ fn exec_pattern() -> Result<(), ()> {
 #[test]
 fn exec_pyimport_test() -> Result<(), ()> {
     // HACK: When running the test with Windows, the exit code is 1 (the cause is unknown)
-    if cfg!(windows) && env_python_version().minor < Some(8) {
+    if cfg!(windows) && env_python_version().unwrap().minor < Some(8) {
         expect_end_with("tests/should_ok/pyimport.er", 2, 1)
     } else {
         expect_success("tests/should_ok/pyimport.er", 2)