From f01b9e4641c4551a3ccf5af4036f0d95cc312f33 Mon Sep 17 00:00:00 2001
From: yum <yum.food.vr@gmail.com>
Date: Tue, 28 Mar 2023 17:35:31 -0700
Subject: Fix virtual env reset

Use `pip freeze` and `pip uninstall` to reset the venv to a near-default
state. Filter out `future` since we need to vendor it. If it ever gets
removed, the installation is borked.
---
 GUI/GUI/GUI/Frame.cpp         | 79 +++++++++++++++++++++++++++++++++++--------
 GUI/GUI/GUI/Frame.h           |  1 +
 GUI/GUI/GUI/PythonWrapper.cpp |  3 ++
 3 files changed, 68 insertions(+), 15 deletions(-)

diff --git a/GUI/GUI/GUI/Frame.cpp b/GUI/GUI/GUI/Frame.cpp
index deccb52..c9b7281 100644
--- a/GUI/GUI/GUI/Frame.cpp
+++ b/GUI/GUI/GUI/Frame.cpp
@@ -370,6 +370,11 @@ Frame::Frame()
 		env_proc_ = p.get_future();
 		p.set_value(true);
 	}
+	{
+		auto p = std::promise<void>();
+		reset_venv_proc_ = p.get_future();
+		p.set_value();
+	}
 
     auto* main_panel = new wxPanel(this, ID_MAIN_PANEL);
 	main_panel_ = main_panel;
@@ -1836,24 +1841,68 @@ void Frame::OnClearPip(wxCommandEvent& event)
 
 void Frame::OnResetVenv(wxCommandEvent& event)
 {
-	Log(debug_out_, "Resetting virtual environment... ");
+    auto status = reset_venv_proc_.wait_for(std::chrono::seconds(0));
+    if (status != std::future_status::ready) {
+		Log(debug_out_, "Virtual environment reset already running\n");
+		return;
+	}
 
-	const std::string py_dir = "Resources/Python/Lib/site-packages";
+    /*
+    Equivalent shell:
+    python -m pip freeze > venv_pkgs.txt
+    python -m pip uninstall -r venv_pkgs.txt
+    rm venv_pkgs.txt
+    */
 
-    if (!std::filesystem::is_directory(py_dir)) {
-		Log(debug_out_, "Python package directory not exist at {}, assuming "
-            "already deleted!\n", py_dir);
-        return;
-    }
+    reset_venv_proc_ = std::move(std::async(std::launch::async, [&]() {
+		Log(debug_out_, "Resetting virtual environment...\n");
 
-	std::error_code err;
-	if (std::filesystem::remove_all(py_dir, err)) {
-		Log(debug_out_, "success!\n");
-	}
-	else {
-		wxLogError("Failed to reset virtual environment: %s", err.message());
-		Log(debug_out_, "failed!\n");
-	}
+        {
+            std::stringstream pkg_list_ss;
+            auto out_cb = [&](const std::string& out, const std::string& err) {
+                Log(debug_out_, "{}", out);
+                Log(debug_out_, "{}", err);
+                pkg_list_ss << out;
+            };
+            auto in_cb = [&](std::string& in) {};
+            Log(debug_out_, "Freezing packages...\n");
+            if (!PythonWrapper::InvokeWithArgs({ "-m pip freeze" }, out_cb, in_cb)) {
+                Log(debug_out_, "failed!\n");
+                return;
+            }
+
+            std::stringstream pkg_list_ss2;
+            std::string pkg_line;
+            while (std::getline(pkg_list_ss, pkg_line)) {
+                if (pkg_line.find("future") != std::string::npos) {
+                    continue;
+                }
+                pkg_list_ss2 << pkg_line << std::endl;
+            }
+
+            std::ofstream pkgs_ofs("venv_pkgs.txt");
+            pkgs_ofs << pkg_list_ss2.str();
+            pkgs_ofs.close();
+        }
+
+        // For now, leave venv_pkgs.txt on disk for better debuggability.
+		//ScopeGuard venv_pkgs_cleanup([]() { std::filesystem::remove("venv_pkgs.txt"); });
+
+		{
+			auto out_cb = [&](const std::string& out, const std::string& err) {
+				Log(debug_out_, "{}", out);
+				Log(debug_out_, "{}", err);
+			};
+			auto in_cb = [&](std::string& in) {};
+			Log(debug_out_, "Uninstalling packages...\n");
+			if (!PythonWrapper::InvokeWithArgs({ "-m pip uninstall -y -r venv_pkgs.txt" }, out_cb, in_cb)) {
+				Log(debug_out_, "failed!\n");
+				return;
+			}
+		}
+
+		Log(debug_out_, "Virtual environment reset done!\n");
+	}));
 }
 
 void Frame::OnClearOSC(wxCommandEvent& event)
diff --git a/GUI/GUI/GUI/Frame.h b/GUI/GUI/GUI/Frame.h
index 98a62c4..8bb88ea 100644
--- a/GUI/GUI/GUI/Frame.h
+++ b/GUI/GUI/GUI/Frame.h
@@ -91,6 +91,7 @@ private:
     std::future<bool> unity_app_;
     std::future<bool> dump_mics_;
     std::future<bool> env_proc_;
+    std::future<void> reset_venv_proc_;
 
     wxTimer py_app_drain_;
 
diff --git a/GUI/GUI/GUI/PythonWrapper.cpp b/GUI/GUI/GUI/PythonWrapper.cpp
index f920a43..ae0fa49 100644
--- a/GUI/GUI/GUI/PythonWrapper.cpp
+++ b/GUI/GUI/GUI/PythonWrapper.cpp
@@ -144,12 +144,14 @@ bool SetAffinityMask(
 			return false;
 		}
 
+#if 0
 		{
 			std::ostringstream oss;
 			oss << "Set affinity mask to " << affinity_mask <<
 				", i.e. processor " << processor_id << std::endl;
 			out_cb(oss.str(), "");
 		}
+#endif
 	}
 	return true;
 }
@@ -470,6 +472,7 @@ bool PythonWrapper::InstallPip(
 	const std::function<bool()>&& run_cb) {
 	std::filesystem::path pip_flag = "Resources/Python/.pip_installed";
 	if (std::filesystem::exists(pip_flag)) {
+		out_cb("Pip flag exists, already installed\n", "");
 		return true;
 	}
 
-- 
cgit v1.2.3