bpo-38133: Allow py.exe launcher to locate installations from the Microsoft Store (GH-16025)

2 files changed, 211 insertions, 101 deletions

diff --git a/Misc/NEWS.d/next/Windows/2019-09-12-12-05-55.bpo-38133.yFeRGS.rst b/Misc/NEWS.d/next/Windows/2019-09-12-12-05-55.bpo-38133.yFeRGS.rst
new file mode 100644
index 0000000..3fbf016
--- /dev/null
+++ b/Misc/NEWS.d/next/Windows/2019-09-12-12-05-55.bpo-38133.yFeRGS.rst
@@ -0,0 +1,2 @@
+Allow py.exe launcher to locate installations from the Microsoft Store and
+improve display of active virtual environments.
diff --git a/PC/launcher.c b/PC/launcher.c
index d23637c..a260860 100644
--- a/PC/launcher.c
+++ b/PC/launcher.c
@@ -162,12 +162,13 @@ static wchar_t * get_env(wchar_t * key)
 #endif
 #endif
 
-#define MAX_VERSION_SIZE    4
+#define MAX_VERSION_SIZE    8
 
 typedef struct {
     wchar_t version[MAX_VERSION_SIZE]; /* m.n */
     int bits;   /* 32 or 64 */
     wchar_t executable[MAX_PATH];
+    wchar_t exe_display[MAX_PATH];
 } INSTALLED_PYTHON;
 
 /*
@@ -185,10 +186,18 @@ static size_t num_installed_pythons = 0;
  * The version name can be longer than MAX_VERSION_SIZE, but will be
  * truncated to just X.Y for comparisons.
  */
-#define IP_BASE_SIZE 40
+#define IP_BASE_SIZE 80
 #define IP_VERSION_SIZE 8
 #define IP_SIZE (IP_BASE_SIZE + IP_VERSION_SIZE)
 #define CORE_PATH L"SOFTWARE\\Python\\PythonCore"
+/*
+ * Installations from the Microsoft Store will set the same registry keys,
+ * but because of a limitation in Windows they cannot be enumerated normally
+ * (unless you have no other Python installations... which is probably false
+ * because that's the most likely way to get this launcher!)
+ * This key is under HKEY_LOCAL_MACHINE
+ */
+#define LOOKASIDE_PATH L"SOFTWARE\\Microsoft\\AppModel\\Lookaside\\user\\Software\\Python\\PythonCore"
 
 static wchar_t * location_checks[] = {
     L"\\",
@@ -201,7 +210,7 @@ static wchar_t * location_checks[] = {
 };
 
 static INSTALLED_PYTHON *
-find_existing_python(wchar_t * path)
+find_existing_python(const wchar_t * path)
 {
     INSTALLED_PYTHON * result = NULL;
     size_t i;
@@ -216,15 +225,32 @@ find_existing_python(wchar_t * path)
     return result;
 }
 
+static INSTALLED_PYTHON *
+find_existing_python2(int bits, const wchar_t * version)
+{
+    INSTALLED_PYTHON * result = NULL;
+    size_t i;
+    INSTALLED_PYTHON * ip;
+
+    for (i = 0, ip = installed_pythons; i < num_installed_pythons; i++, ip++) {
+        if (bits == ip->bits && _wcsicmp(version, ip->version) == 0) {
+            result = ip;
+            break;
+        }
+    }
+    return result;
+}
+
 static void
-locate_pythons_for_key(HKEY root, REGSAM flags)
+_locate_pythons_for_key(HKEY root, LPCWSTR subkey, REGSAM flags, int bits,
+                        int display_name_only)
 {
     HKEY core_root, ip_key;
-    LSTATUS status = RegOpenKeyExW(root, CORE_PATH, 0, flags, &core_root);
+    LSTATUS status = RegOpenKeyExW(root, subkey, 0, flags, &core_root);
     wchar_t message[MSGSIZE];
     DWORD i;
     size_t n;
-    BOOL ok;
+    BOOL ok, append_name;
     DWORD type, data_size, attrs;
     INSTALLED_PYTHON * ip, * pip;
     wchar_t ip_version[IP_VERSION_SIZE];
@@ -252,8 +278,16 @@ locate_pythons_for_key(HKEY root, REGSAM flags)
             else {
                 wcsncpy_s(ip->version, MAX_VERSION_SIZE, ip_version,
                           MAX_VERSION_SIZE-1);
+                /* Still treating version as "x.y" rather than sys.winver
+                 * When PEP 514 tags are properly used, we shouldn't need
+                 * to strip this off here.
+                 */
+                check = wcsrchr(ip->version, L'-');
+                if (check && !wcscmp(check, L"-32")) {
+                    *check = L'\0';
+                }
                 _snwprintf_s(ip_path, IP_SIZE, _TRUNCATE,
-                             L"%ls\\%ls\\InstallPath", CORE_PATH, ip_version);
+                             L"%ls\\%ls\\InstallPath", subkey, ip_version);
                 status = RegOpenKeyExW(root, ip_path, 0, flags, &ip_key);
                 if (status != ERROR_SUCCESS) {
                     winerror(status, message, MSGSIZE);
@@ -262,42 +296,57 @@ locate_pythons_for_key(HKEY root, REGSAM flags)
                     continue;
                 }
                 data_size = sizeof(ip->executable) - 1;
-                status = RegQueryValueExW(ip_key, NULL, NULL, &type,
+                append_name = FALSE;
+                status = RegQueryValueExW(ip_key, L"ExecutablePath", NULL, &type,
                                           (LPBYTE)ip->executable, &data_size);
+                if (status != ERROR_SUCCESS || type != REG_SZ || !data_size) {
+                    append_name = TRUE;
+                    data_size = sizeof(ip->executable) - 1;
+                    status = RegQueryValueExW(ip_key, NULL, NULL, &type,
+                                              (LPBYTE)ip->executable, &data_size);
+                    if (status != ERROR_SUCCESS) {
+                        winerror(status, message, MSGSIZE);
+                        debug(L"%ls\\%ls: %ls\n", key_name, ip_path, message);
+                        RegCloseKey(ip_key);
+                        continue;
+                    }
+                }
                 RegCloseKey(ip_key);
-                if (status != ERROR_SUCCESS) {
-                    winerror(status, message, MSGSIZE);
-                    debug(L"%ls\\%ls: %ls\n", key_name, ip_path, message);
+                if (type != REG_SZ) {
                     continue;
                 }
-                if (type == REG_SZ) {
-                    data_size = data_size / sizeof(wchar_t) - 1;  /* for NUL */
-                    if (ip->executable[data_size - 1] == L'\\')
-                        --data_size; /* reg value ended in a backslash */
-                    /* ip->executable is data_size long */
-                    for (checkp = location_checks; *checkp; ++checkp) {
-                        check = *checkp;
+
+                data_size = data_size / sizeof(wchar_t) - 1;  /* for NUL */
+                if (ip->executable[data_size - 1] == L'\\')
+                    --data_size; /* reg value ended in a backslash */
+                /* ip->executable is data_size long */
+                for (checkp = location_checks; *checkp; ++checkp) {
+                    check = *checkp;
+                    if (append_name) {
                         _snwprintf_s(&ip->executable[data_size],
                                      MAX_PATH - data_size,
                                      MAX_PATH - data_size,
                                      L"%ls%ls", check, PYTHON_EXECUTABLE);
-                        attrs = GetFileAttributesW(ip->executable);
-                        if (attrs == INVALID_FILE_ATTRIBUTES) {
-                            winerror(GetLastError(), message, MSGSIZE);
-                            debug(L"locate_pythons_for_key: %ls: %ls",
-                                  ip->executable, message);
-                        }
-                        else if (attrs & FILE_ATTRIBUTE_DIRECTORY) {
-                            debug(L"locate_pythons_for_key: '%ls' is a \
-directory\n",
-                                  ip->executable, attrs);
-                        }
-                        else if (find_existing_python(ip->executable)) {
-                            debug(L"locate_pythons_for_key: %ls: already \
-found\n", ip->executable);
-                        }
-                        else {
-                            /* check the executable type. */
+                    }
+                    attrs = GetFileAttributesW(ip->executable);
+                    if (attrs == INVALID_FILE_ATTRIBUTES) {
+                        winerror(GetLastError(), message, MSGSIZE);
+                        debug(L"locate_pythons_for_key: %ls: %ls",
+                              ip->executable, message);
+                    }
+                    else if (attrs & FILE_ATTRIBUTE_DIRECTORY) {
+                        debug(L"locate_pythons_for_key: '%ls' is a directory\n",
+                              ip->executable, attrs);
+                    }
+                    else if (find_existing_python(ip->executable)) {
+                        debug(L"locate_pythons_for_key: %ls: already found\n",
+                              ip->executable);
+                    }
+                    else {
+                        /* check the executable type. */
+                        if (bits) {
+                            ip->bits = bits;
+                        } else {
                             ok = GetBinaryTypeW(ip->executable, &attrs);
                             if (!ok) {
                                 debug(L"Failure getting binary type: %ls\n",
@@ -310,32 +359,48 @@ found\n", ip->executable);
                                     ip->bits = 32;
                                 else
                                     ip->bits = 0;
-                                if (ip->bits == 0) {
-                                    debug(L"locate_pythons_for_key: %ls: \
+                            }
+                        }
+                        if (ip->bits == 0) {
+                            debug(L"locate_pythons_for_key: %ls: \
 invalid binary type: %X\n",
-                                          ip->executable, attrs);
+                                  ip->executable, attrs);
+                        }
+                        else {
+                            if (display_name_only) {
+                                /* display just the executable name. This is
+                                 * primarily for the Store installs */
+                                const wchar_t *name = wcsrchr(ip->executable, L'\\');
+                                if (name) {
+                                    wcscpy_s(ip->exe_display, MAX_PATH, name+1);
                                 }
-                                else {
-                                    if (wcschr(ip->executable, L' ') != NULL) {
-                                        /* has spaces, so quote */
-                                        n = wcslen(ip->executable);
-                                        memmove(&ip->executable[1],
-                                                ip->executable, n * sizeof(wchar_t));
-                                        ip->executable[0] = L'\"';
-                                        ip->executable[n + 1] = L'\"';
-                                        ip->executable[n + 2] = L'\0';
-                                    }
-                                    debug(L"locate_pythons_for_key: %ls \
-is a %dbit executable\n",
-                                        ip->executable, ip->bits);
-                                    ++num_installed_pythons;
-                                    pip = ip++;
-                                    if (num_installed_pythons >=
-                                        MAX_INSTALLED_PYTHONS)
-                                        break;
-                                    /* Copy over the attributes for the next */
-                                    *ip = *pip;
+                            }
+                            if (wcschr(ip->executable, L' ') != NULL) {
+                                /* has spaces, so quote, and set original as
+                                 * the display name */
+                                if (!ip->exe_display[0]) {
+                                    wcscpy_s(ip->exe_display, MAX_PATH, ip->executable);
                                 }
+                                n = wcslen(ip->executable);
+                                memmove(&ip->executable[1],
+                                        ip->executable, n * sizeof(wchar_t));
+                                ip->executable[0] = L'\"';
+                                ip->executable[n + 1] = L'\"';
+                                ip->executable[n + 2] = L'\0';
+                            }
+                            debug(L"locate_pythons_for_key: %ls \
+is a %dbit executable\n",
+                                ip->executable, ip->bits);
+                            if (find_existing_python2(ip->bits, ip->version)) {
+                                debug(L"locate_pythons_for_key: %ls-%i: already \
+found\n", ip->version, ip->bits);
+                            }
+                            else {
+                                ++num_installed_pythons;
+                                pip = ip++;
+                                if (num_installed_pythons >=
+                                    MAX_INSTALLED_PYTHONS)
+                                    break;
                             }
                         }
                     }
@@ -360,8 +425,62 @@ compare_pythons(const void * p1, const void * p2)
 }
 
 static void
+locate_pythons_for_key(HKEY root, REGSAM flags)
+{
+    _locate_pythons_for_key(root, CORE_PATH, flags, 0, FALSE);
+}
+
+static void
+locate_store_pythons()
+{
+#if defined(_M_X64)
+    /* 64bit process, so look in native registry */
+    _locate_pythons_for_key(HKEY_LOCAL_MACHINE, LOOKASIDE_PATH,
+                            KEY_READ, 64, TRUE);
+#else
+    /* 32bit process, so check that we're on 64bit OS */
+    BOOL f64 = FALSE;
+    if (IsWow64Process(GetCurrentProcess(), &f64) && f64) {
+        _locate_pythons_for_key(HKEY_LOCAL_MACHINE, LOOKASIDE_PATH,
+                                KEY_READ | KEY_WOW64_64KEY, 64, TRUE);
+    }
+#endif
+}
+
+static void
+locate_venv_python()
+{
+    static wchar_t venv_python[MAX_PATH];
+    INSTALLED_PYTHON * ip;
+    wchar_t *virtual_env = get_env(L"VIRTUAL_ENV");
+    DWORD attrs;
+
+    /* Check for VIRTUAL_ENV environment variable */
+    if (virtual_env == NULL || virtual_env[0] == L'\0') {
+        return;
+    }
+
+    /* Check for a python executable in the venv */
+    debug(L"Checking for Python executable in virtual env '%ls'\n", virtual_env);
+    _snwprintf_s(venv_python, MAX_PATH, _TRUNCATE,
+            L"%ls\\Scripts\\%ls", virtual_env, PYTHON_EXECUTABLE);
+    attrs = GetFileAttributesW(venv_python);
+    if (attrs == INVALID_FILE_ATTRIBUTES) {
+        debug(L"Python executable %ls missing from virtual env\n", venv_python);
+        return;
+    }
+
+    ip = &installed_pythons[num_installed_pythons++];
+    wcscpy_s(ip->executable, MAX_PATH, venv_python);
+    ip->bits = 0;
+    wcscpy_s(ip->version, MAX_VERSION_SIZE, L"venv");
+}
+
+static void
 locate_all_pythons()
 {
+    /* venv Python is highest priority */
+    locate_venv_python();
 #if defined(_M_X64)
     /* If we are a 64bit process, first hit the 32bit keys. */
     debug(L"locating Pythons in 32bit registry\n");
@@ -380,6 +499,8 @@ locate_all_pythons()
     debug(L"locating Pythons in native registry\n");
     locate_pythons_for_key(HKEY_CURRENT_USER, KEY_READ);
     locate_pythons_for_key(HKEY_LOCAL_MACHINE, KEY_READ);
+    /* Store-installed Python is lowest priority */
+    locate_store_pythons();
     qsort(installed_pythons, num_installed_pythons, sizeof(INSTALLED_PYTHON),
           compare_pythons);
 }
@@ -416,31 +537,6 @@ find_python_by_version(wchar_t const * wanted_ver)
 }
 
 
-static wchar_t *
-find_python_by_venv()
-{
-    static wchar_t venv_python[MAX_PATH];
-    wchar_t *virtual_env = get_env(L"VIRTUAL_ENV");
-    DWORD attrs;
-
-    /* Check for VIRTUAL_ENV environment variable */
-    if (virtual_env == NULL || virtual_env[0] == L'\0') {
-        return NULL;
-    }
-
-    /* Check for a python executable in the venv */
-    debug(L"Checking for Python executable in virtual env '%ls'\n", virtual_env);
-    _snwprintf_s(venv_python, MAX_PATH, _TRUNCATE,
-            L"%ls\\Scripts\\%ls", virtual_env, PYTHON_EXECUTABLE);
-    attrs = GetFileAttributesW(venv_python);
-    if (attrs == INVALID_FILE_ATTRIBUTES) {
-        debug(L"Python executable %ls missing from virtual env\n", venv_python);
-        return NULL;
-    }
-
-    return venv_python;
-}
-
 static wchar_t appdata_ini_path[MAX_PATH];
 static wchar_t launcher_ini_path[MAX_PATH];
 
@@ -523,9 +619,12 @@ locate_python(wchar_t * wanted_ver, BOOL from_shebang)
     }
     else {
         *last_char = L'\0'; /* look for an overall default */
-        configured_value = get_configured_value(config_key);
-        if (configured_value)
-            result = find_python_by_version(configured_value);
+        result = find_python_by_version(L"venv");
+        if (result == NULL) {
+            configured_value = get_configured_value(config_key);
+            if (configured_value)
+                result = find_python_by_version(configured_value);
+        }
         /* Not found a value yet - try by major version.
          * If we're looking for an interpreter specified in a shebang line,
          * we want to try Python 2 first, then Python 3 (for Unix and backward
@@ -1443,7 +1542,8 @@ show_python_list(wchar_t ** argv)
     INSTALLED_PYTHON * defpy = locate_python(L"", FALSE);
     size_t i = 0;
     wchar_t *p = argv[1];
-    wchar_t *fmt = L"\n -%ls-%d"; /* print VER-BITS */
+    wchar_t *ver_fmt = L"-%ls-%d";
+    wchar_t *fmt = L"\n %ls";
     wchar_t *defind = L" *"; /* Default indicator */
 
     /*
@@ -1452,8 +1552,8 @@ show_python_list(wchar_t ** argv)
     */
     fwprintf(stderr,
              L"Installed Pythons found by %s Launcher for Windows", argv[0]);
-    if (!_wcsicmp(p, L"-0p") || !_wcsicmp(p, L"--list-paths")) /* Show path? */
-        fmt = L"\n -%ls-%d\t%ls"; /* print VER-BITS path */
+    if (!_wcsicmp(p, L"-0p") || !_wcsicmp(p, L"--list-paths"))
+        fmt = L"\n %-15ls%ls"; /* include path */
 
     if (num_installed_pythons == 0) /* We have somehow got here without searching for pythons */
         locate_all_pythons(); /* Find them, Populates installed_pythons */
@@ -1463,9 +1563,22 @@ show_python_list(wchar_t ** argv)
     else
     {
         for (i = 0; i < num_installed_pythons; i++, ip++) {
-            fwprintf(stdout, fmt, ip->version, ip->bits, ip->executable);
+            wchar_t version[BUFSIZ];
+            if (wcscmp(ip->version, L"venv") == 0) {
+                wcscpy_s(version, BUFSIZ, L"(venv)");
+            }
+            else {
+                swprintf_s(version, BUFSIZ, ver_fmt, ip->version, ip->bits);
+            }
+
+            if (ip->exe_display[0]) {
+                fwprintf(stdout, fmt, version, ip->exe_display);
+            }
+            else {
+                fwprintf(stdout, fmt, version, ip->executable);
+            }
             /* If there is a default indicate it */
-            if ((defpy != NULL) && !_wcsicmp(ip->executable, defpy->executable))
+            if (defpy == ip)
                 fwprintf(stderr, defind);
         }
     }
@@ -1850,16 +1963,11 @@ installed, use -0 for available pythons", &p[1]);
             executable = NULL; /* Info call only */
         }
         else {
-            /* Look for an active virtualenv */
-            executable = find_python_by_venv();
-
-            /* If we didn't find one, look for the default Python */
-            if (executable == NULL) {
-                ip = locate_python(L"", FALSE);
-                if (ip == NULL)
-                    error(RC_NO_PYTHON, L"Can't find a default Python.");
-                executable = ip->executable;
-            }
+            /* look for the default Python */
+            ip = locate_python(L"", FALSE);
+            if (ip == NULL)
+                error(RC_NO_PYTHON, L"Can't find a default Python.");
+            executable = ip->executable;
         }
     }
     if (executable != NULL)