Ensure Start shortcuts and file associations use good working directory. (#136)

Fixes #134

Author: zooba

_msbuild.py

@@ -78,6 +78,7 @@ class ResourceFile(CSourceFile):
     CFunction('file_locked_delete'),
     CFunction('package_get_root'),
     CFunction('shortcut_create'),
+    CFunction('shortcut_default_cwd'),
     CFunction('shortcut_get_start_programs'),
     CFunction('hide_file'),
     CFunction('fd_supports_vt100'),
@@ -97,7 +98,11 @@ def main_exe(name):
         VersionInfo(FileDescription="Python Install Manager"),
         CPP_SETTINGS,
         ItemDefinition('ClCompile', PreprocessorDefinitions=Prepend(f'EXE_NAME=L"{name}";')),
-        ItemDefinition('Link', SubSystem='CONSOLE', DelayLoadDLLs=f"{DLL_NAME}.dll"),
+        ItemDefinition('Link',
+            SubSystem='CONSOLE',
+            DelayLoadDLLs=f'{DLL_NAME}.dll;ole32.dll;shell32.dll;advapi32.dll',
+            DisableSpecificWarnings=Prepend('4199;'),
+        ),
         INCLUDE_TMPDIR,
         Manifest('default.manifest'),
         ResourceFile('pyicon.rc'),
@@ -115,7 +120,11 @@ def mainw_exe(name):
     return CProject(name,
         VersionInfo(FileDescription="Python Install Manager (windowed)"),
         CPP_SETTINGS,
-        ItemDefinition('Link', SubSystem='WINDOWS', DelayLoadDLLs=f"{DLL_NAME}.dll"),
+        ItemDefinition('Link',
+            SubSystem='WINDOWS',
+            DelayLoadDLLs=f'{DLL_NAME}.dll;ole32.dll;shell32.dll;advapi32.dll',
+            DisableSpecificWarnings=Prepend('4199;'),
+        ),
         INCLUDE_TMPDIR,
         ItemDefinition('ClCompile', PreprocessorDefinitions=Prepend(f'EXE_NAME=L"{name}";')),
         ItemDefinition('ClCompile', PreprocessorDefinitions=Prepend("PY_WINDOWED=1;")),

_msbuild_test.py

@@ -44,6 +44,7 @@
         CFunction('file_locked_delete'),
         CFunction('package_get_root'),
         CFunction('shortcut_create'),
+        CFunction('shortcut_default_cwd'),
         CFunction('shortcut_get_start_programs'),
         CFunction('hide_file'),
         CFunction('fd_supports_vt100'),

src/_native/shortcut.cpp

@@ -2,6 +2,7 @@
 #include <windows.h>
 #include <shlobj.h>
 #include <shlguid.h>
+#include <iterator>
 #include "helpers.h"
 
 extern "C" {
@@ -111,6 +112,27 @@ shortcut_get_start_programs(PyObject *, PyObject *, PyObject *)
 }
 
 
+PyObject *
+shortcut_default_cwd(PyObject *, PyObject *, PyObject *)
+{
+    static const KNOWNFOLDERID fids[] = { FOLDERID_Documents, FOLDERID_Profile };
+    for (auto i = std::begin(fids); i != std::end(fids); ++i) {
+        wchar_t *path;
+        if (SUCCEEDED(SHGetKnownFolderPath(
+            *i,
+            KF_FLAG_NO_PACKAGE_REDIRECTION,
+            NULL,
+            &path
+        ))) {
+            PyObject *r = PyUnicode_FromWideChar(path, -1);
+            CoTaskMemFree(path);
+            return r;
+        }
+    }
+    return Py_GetConstant(Py_CONSTANT_NONE);
+}
+
+
 PyObject *
 hide_file(PyObject *, PyObject *args, PyObject *kwargs)
 {

src/manage/commands.py

@@ -39,6 +39,8 @@
 
 
 WELCOME = f"""!B!Python install manager was successfully updated to {__version__}.!W!
+!Y!Start menu shortcuts have been changed in this update.!W!
+Run !G!py install --refresh!W! to update any existing shortcuts.
 """
 
 # The 'py help' or 'pymanager help' output is constructed by these default docs,

src/manage/startutils.py

@@ -51,7 +51,8 @@ def _make(root, prefix, item, allow_warn=True):
         lnk,
         target,
         arguments=_unprefix(item.get("Arguments"), prefix),
-        working_directory=_unprefix(item.get("WorkingDirectory"), prefix),
+        working_directory=_unprefix(item.get("WorkingDirectory"), prefix)
+            or _native.shortcut_default_cwd(),
         icon=_unprefix(item.get("Icon"), prefix),
         icon_index=item.get("IconIndex", 0),
     )

src/pymanager/main.cpp

@@ -2,6 +2,7 @@
 #include <string.h>
 
 #include <Windows.h>
+#include <shlwapi.h>
 #include <stdio.h>
 
 #include <shellapi.h>
@@ -23,6 +24,11 @@
 #define PY_WINDOWED 0
 #endif
 
+// Uncomment to add the --__fix-cwd private argument, which will reset the
+// working directory to the script location or user documents/profile before
+// running Python.
+//#define ENABLE_FIX_CWD
+
 struct {
     PyObject *mod;
     PyObject *no_install_found_error;
@@ -483,6 +489,38 @@ locate_runtime(
 }
 
 
+#ifdef ENABLE_FIX_CWD
+static int
+fix_working_directory(const std::wstring &script)
+{
+    HRESULT hr;
+    // If we have a script, use its parent directory
+    if (!script.empty()) {
+        auto end = script.find_last_of(L"/\\");
+        if (end != script.npos) {
+            std::wstring current_dir(script.data(), end);
+            SetCurrentDirectoryW(current_dir.c_str());
+            return 0;
+        }
+    }
+    // If we have no script, assume the user's documents folder
+    wchar_t *path;
+    if (SUCCEEDED(hr = SHGetKnownFolderPath(FOLDERID_Documents, 0, NULL, &path))) {
+        SetCurrentDirectoryW(path);
+        CoTaskMemFree(path);
+        return 0;
+    }
+    // As a fallback, use the user's profile (e.g. for SYSTEM)
+    if (SUCCEEDED(hr = SHGetKnownFolderPath(FOLDERID_Profile, 0, NULL, &path))) {
+        SetCurrentDirectoryW(path);
+        CoTaskMemFree(path);
+        return 0;
+    }
+    return hr;
+}
+#endif
+
+
 int
 wmain(int argc, wchar_t **argv)
 {
@@ -504,6 +542,9 @@ wmain(int argc, wchar_t **argv)
 
     const wchar_t *default_cmd;
     bool use_commands, use_cli_tag, use_shebangs, use_autoinstall;
+    #ifdef ENABLE_FIX_CWD
+    bool fix_cwd = false;
+    #endif
     per_exe_settings(argc, argv, &default_cmd, &use_commands, &use_cli_tag, &use_shebangs, &use_autoinstall);
 
     if (use_commands) {
@@ -519,6 +560,12 @@ wmain(int argc, wchar_t **argv)
         // We handle 'exec' in native code, so it won't be in the above list
         if (!wcscmp(argv[1], L"exec")) {
             skip_argc += 1;
+            #ifdef ENABLE_FIX_CWD
+            if (!wcscmp(argv[2], L"--__fix-cwd")) {
+                fix_cwd = true;
+                skip_argc += 1;
+            }
+            #endif
             use_cli_tag = argc >= 3;
             use_shebangs = argc >= 3;
             default_cmd = NULL;
@@ -570,6 +617,15 @@ wmain(int argc, wchar_t **argv)
     // Theoretically shouldn't matter, but might help reduce memory usage.
     close_python();
 
+    #ifdef ENABLE_FIX_CWD
+    if (fix_cwd) {
+        err = fix_working_directory(script);
+        if (err) {
+            fprintf(stderr, "[WARN] Failed to fix working directory (0x%08X).\n", err);
+        }
+    }
+    #endif
+
     err = launch(executable.c_str(), args.c_str(), skip_argc, &exitCode);
 
     // TODO: Consider sharing print_error() with launcher.cpp