From 760d8e74e15b81276d7b3a39d168036d1f7aa141 Mon Sep 17 00:00:00 2001
From: Dustin Spicuzza <dustin@virtualroadside.com>
Date: Wed, 12 Feb 2025 00:18:36 -0500
Subject: [PATCH] Add Python compatibility check

- Fixes #125
---
 robotpy_installer/cli_deploy.py    | 54 +++++++++++++++++++++++++++---
 robotpy_installer/installer.py     |  3 +-
 robotpy_installer/roborio_utils.py | 15 +++++++++
 3 files changed, 66 insertions(+), 6 deletions(-)

diff --git a/robotpy_installer/cli_deploy.py b/robotpy_installer/cli_deploy.py
index 35a1f27..a1a98de 100644
--- a/robotpy_installer/cli_deploy.py
+++ b/robotpy_installer/cli_deploy.py
@@ -17,6 +17,7 @@
 
 from . import pypackages, pyproject, roborio_utils, sshcontroller
 from .installer import PipInstallError, PythonMissingError, RobotpyInstaller
+from .installer import _ROBOTPY_PYTHON_VERSION_TUPLE as required_pyversion
 from .errors import Error
 from .utils import handle_cli_error, print_err, yesno
 
@@ -349,6 +350,14 @@ def _get_robot_packages(
             self._robot_packages = pypackages.make_packages(rio_packages)
         return self._robot_packages
 
+    def _clear_pip_packages(self, installer: RobotpyInstaller):
+        rio_packages = self._get_robot_packages(installer.ssh)
+        to_uninstall = [p for p in rio_packages.keys() if p != "pip"]
+        if to_uninstall:
+            installer.pip_uninstall(to_uninstall)
+
+        self._packages_in_cache = None
+
     def _ensure_requirements(
         self,
         project: typing.Optional[pyproject.RobotPyProjectToml],
@@ -361,6 +370,7 @@ def _ensure_requirements(
         no_uninstall: bool,
     ):
         python_exists = False
+        python_invalid: typing.Union[bool, str] = False
         requirements_installed = False
 
         installer = RobotpyInstaller()
@@ -383,6 +393,33 @@ def _ensure_requirements(
             if not python_exists:
                 logger.warning("Python is not installed on RoboRIO")
 
+        if python_exists:
+            with wrap_ssh_error("getting python version"):
+                python_version = roborio_utils.get_python3_version(ssh)
+
+            if python_version != required_pyversion:
+                python_exists = False
+                m, mn = python_version
+                rm, rmn = required_pyversion
+                python_invalid = f"python{m}{mn}"
+
+                if no_install:
+                    raise Error(
+                        f"Unsupported version of python ({m}.{mn}) was found on the roboRIO\n"
+                        "- could not update it because no-install was specified\n"
+                    )
+
+                # Warn the user before changing their rio
+                print(
+                    "\n"
+                    f"Deployer has detected that the version of Python installed on the RoboRIO ({m}.{mn})\n"
+                    "is not supported by this installer. The installer will now uninstall that\n"
+                    f"and install Python {rm}.{rmn}.\n"
+                )
+
+                if not yesno("Reinstall Python"):
+                    raise Error("User declined reinstallation")
+
         if python_exists:
             if no_install:
                 requirements_installed = True
@@ -440,6 +477,7 @@ def _ensure_requirements(
         if (
             cpp_java_exists
             or not python_exists
+            or python_invalid
             or not requirements_installed
             or not kill_script_updated
         ):
@@ -469,6 +507,16 @@ def _ensure_requirements(
                 if cpp_java_exists:
                     roborio_utils.uninstall_cpp_java_admin(installer.ssh)
 
+                if python_invalid:
+                    with wrap_ssh_error("uninstalling python"):
+                        self._clear_pip_packages(installer)
+                        logger.info("Uninstalling %s from RoboRIO", python_invalid)
+                        installer.ssh.exec_cmd(
+                            f"opkg remove {python_invalid}",
+                            check=True,
+                            print_output=True,
+                        )
+
                 if not python_exists:
                     try:
                         installer.install_python()
@@ -508,11 +556,7 @@ def _ensure_requirements(
                         # environment is to first clear the environment.
                         # - can't do a partial uninstall without completely
                         #   resolving everything
-
-                        rio_packages = self._get_robot_packages(installer.ssh)
-                        to_uninstall = [p for p in rio_packages.keys() if p != "pip"]
-                        if to_uninstall:
-                            installer.pip_uninstall(to_uninstall)
+                        self._clear_pip_packages(installer)
 
                     logger.info("Installing project requirements on RoboRIO:")
                     for package in packages:
diff --git a/robotpy_installer/installer.py b/robotpy_installer/installer.py
index 8272448..e324f40 100755
--- a/robotpy_installer/installer.py
+++ b/robotpy_installer/installer.py
@@ -36,7 +36,8 @@
 ]
 
 _ROBOTPY_PYTHON_PLATFORM = "linux_roborio"
-_ROBOTPY_PYTHON_VERSION_NUM = "313"
+_ROBOTPY_PYTHON_VERSION_TUPLE = (3, 13)
+_ROBOTPY_PYTHON_VERSION_NUM = "".join(map(str, _ROBOTPY_PYTHON_VERSION_TUPLE))
 _ROBOTPY_PYTHON_VERSION = f"python{_ROBOTPY_PYTHON_VERSION_NUM}"
 
 _PIP_STUB_PATH = "/home/admin/rpip"
diff --git a/robotpy_installer/roborio_utils.py b/robotpy_installer/roborio_utils.py
index c31305f..6824dce 100644
--- a/robotpy_installer/roborio_utils.py
+++ b/robotpy_installer/roborio_utils.py
@@ -22,6 +22,21 @@
 kill_script_content: typing.Optional[bytes] = None
 
 
+def get_python3_version(ssh: SshController) -> typing.Tuple[int, int]:
+    r = ssh.check_output(
+        "/usr/local/bin/python3 -c 'import json, sys; json.dump(tuple(sys.version_info), sys.stderr)'"
+    )
+
+    python_version = json.loads(r)
+    assert isinstance(python_version, list)
+    python_version = tuple(python_version[:2])
+    assert len(python_version) == 2
+
+    logger.debug("RoboRIO has Python %s.%s installed", *python_version)
+
+    return python_version
+
+
 def uninstall_cpp_java_lvuser(ssh: SshController) -> bool:
     """
     Frees up disk space by removing FRC C++/Java programs. This runs as lvuser or admin.