From ec70720075fe40e3c15b231725cd0f41ccff0077 Mon Sep 17 00:00:00 2001
From: Axel Soll <axel@soll.xyz>
Date: Wed, 5 May 2021 16:23:55 +0200
Subject: [PATCH] sudo option added to venv pip commands (#104)

* Fix pip virtualenv discovery for softwareupdates

* Use setup.py frm latest OctoPrint version.

* Enable pip install if using sudo.
---
 src/octoprint/util/pip.py                  |  2 +-
 src/octoprint/util/piptestballoon/setup.py | 54 +++++++++++++---------
 2 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/src/octoprint/util/pip.py b/src/octoprint/util/pip.py
index 6de4213ee9..2abafb53c8 100644
--- a/src/octoprint/util/pip.py
+++ b/src/octoprint/util/pip.py
@@ -202,7 +202,7 @@ def _setup_pip(self):
 		# in a virtual environment may we proceed with the --user parameter.
 
 		ok, pip_user, pip_virtual_env, pip_install_dir = self._check_pip_setup(pip_command)
-		if not ok:
+		if not ok and (not pip_sudo or not pip_install_dir):
 			self._logger.error("Cannot use pip")
 			return
 
diff --git a/src/octoprint/util/piptestballoon/setup.py b/src/octoprint/util/piptestballoon/setup.py
index 13dfc2ffdc..036d9a56f4 100644
--- a/src/octoprint/util/piptestballoon/setup.py
+++ b/src/octoprint/util/piptestballoon/setup.py
@@ -1,9 +1,8 @@
-from __future__ import absolute_import, division, print_function
-import os
-import sys
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import, division, print_function, unicode_literals
 
 """
-This "python package" doesn't actually install. This is intenional. It is merely
+This "python package" doesn't actually install. This is intentional. It is merely
 used to figure out some information about the environment a specific pip call
 is running under (installation dir, whether it belongs to a virtual environment,
 whether the install location is writable by the current user), and for that it
@@ -17,32 +16,45 @@
 --verbose or --log flags.
 """
 
+import io
+import os
+import sys
+
+
 def produce_output(stream):
-	from distutils.command.install import install as cmd_install
-	from distutils.dist import Distribution
+    from distutils.command.install import install as cmd_install
+    from distutils.dist import Distribution
+
+    cmd = cmd_install(Distribution())
+    cmd.finalize_options()
+
+    install_dir = cmd.install_lib
+    virtual_env = hasattr(sys, "real_prefix") or (
+        hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix
+    )
+    writable = os.access(install_dir, os.W_OK)
 
-	cmd = cmd_install(Distribution())
-	cmd.finalize_options()
+    lines = [
+        "PIP_INSTALL_DIR={}".format(install_dir),
+        "PIP_VIRTUAL_ENV={}".format(virtual_env),
+        "PIP_WRITABLE={}".format(writable),
+    ]
 
-	install_dir = cmd.install_lib
-	virtual_env = hasattr(sys, "real_prefix")
-	writable = os.access(install_dir, os.W_OK)
+    for line in lines:
+        print(line, file=stream)
 
-	print("PIP_INSTALL_DIR={}".format(install_dir), file=stream)
-	print("PIP_VIRTUAL_ENV={}".format(virtual_env), file=stream)
-	print("PIP_WRITABLE={}".format(writable), file=stream)
-	stream.flush()
+    stream.flush()
 
 
 path = os.environ.get("TESTBALLOON_OUTPUT", None)
 if path is not None:
-	# environment variable set, write to a log
-	path = os.path.abspath(path)
-	with open(path, mode="w+b") as output:
-		produce_output(output)
+    # environment variable set, write to a log
+    path = os.path.abspath(path)
+    with io.open(path, "wt+", encoding="utf-8") as output:
+        produce_output(output)
 else:
-	# write to stdout
-	produce_output(sys.stdout)
+    # write to stdout
+    produce_output(sys.stdout)
 
 # fail intentionally
 sys.exit(-1)