From ea5d04e12919963f5faa611877efe94b3943cab1 Mon Sep 17 00:00:00 2001 From: Steve Dower Date: Fri, 2 May 2025 14:34:53 +0100 Subject: [PATCH] Adds test harness and tests for pyshellext. Fixes #54 --- _msbuild_test.py | 19 +++ src/pyshellext/shellext.cpp | 71 ++++++--- src/pyshellext/shellext.h | 23 +++ src/pyshellext/shellext_test.cpp | 241 +++++++++++++++++++++++++++++++ tests/conftest.py | 46 ++++++ tests/test_pep514utils.py | 45 +----- tests/test_shellext.py | 103 +++++++++++++ 7 files changed, 489 insertions(+), 59 deletions(-) create mode 100644 src/pyshellext/shellext.h create mode 100644 src/pyshellext/shellext_test.cpp create mode 100644 tests/test_shellext.py diff --git a/_msbuild_test.py b/_msbuild_test.py index dfe1820..b5d5c26 100644 --- a/_msbuild_test.py +++ b/_msbuild_test.py @@ -52,4 +52,23 @@ CFunction('reg_rename_key'), source='src/_native', ), + DllPackage('_shellext_test', + PyFile('_native/__init__.py'), + ItemDefinition('ClCompile', + PreprocessorDefinitions=Prepend("PYSHELLEXT_TEST=1;"), + LanguageStandard='stdcpp20', + ), + ItemDefinition('Link', AdditionalDependencies=Prepend("RuntimeObject.lib;")), + CSourceFile('pyshellext/shellext.cpp'), + CSourceFile('pyshellext/shellext_test.cpp'), + IncludeFile('pyshellext/shellext.h'), + CSourceFile('_native/helpers.cpp'), + IncludeFile('_native/helpers.h'), + CFunction('shellext_RegReadStr'), + CFunction('shellext_ReadIdleInstalls'), + CFunction('shellext_ReadAllIdleInstalls'), + CFunction('shellext_PassthroughTitle'), + CFunction('shellext_IdleCommand'), + source='src', + ) ) diff --git a/src/pyshellext/shellext.cpp b/src/pyshellext/shellext.cpp index 4cbc64b..b38f669 100644 --- a/src/pyshellext/shellext.cpp +++ b/src/pyshellext/shellext.cpp @@ -8,13 +8,7 @@ using namespace Microsoft::WRL; -#include -#include -#include -#include - -#include -#include +#include "shellext.h" static HINSTANCE hModule; @@ -23,14 +17,7 @@ static HINSTANCE hModule; #define CLSID_COMMAND_ENUMERATOR "{F82C8CD5-A69C-45CC-ADC6-87FC5F4A7429}" -struct IdleData { - std::wstring title; - std::wstring exe; - std::wstring idle; -}; - - -static LRESULT RegReadStr(HKEY key, LPCWSTR valueName, std::wstring& result) +LRESULT RegReadStr(HKEY key, LPCWSTR valueName, std::wstring& result) { DWORD reg_type; while (true) { @@ -59,7 +46,7 @@ static LRESULT RegReadStr(HKEY key, LPCWSTR valueName, std::wstring& result) } -static HRESULT ReadIdleInstalls(std::vector &idles, HKEY hkPython, LPCWSTR company, REGSAM flags) +HRESULT ReadIdleInstalls(std::vector &idles, HKEY hkPython, LPCWSTR company, REGSAM flags) { HKEY hkCompany = NULL, hkTag = NULL, hkInstall = NULL; LSTATUS err = RegOpenKeyExW( @@ -121,6 +108,8 @@ static HRESULT ReadIdleInstalls(std::vector &idles, HKEY hkPython, LPC data.idle += L"Lib\\idlelib\\idle.pyw"; } } + } else { + err = 0; } } if (err) { @@ -154,11 +143,11 @@ static HRESULT ReadIdleInstalls(std::vector &idles, HKEY hkPython, LPC return S_OK; } -static HRESULT ReadAllIdleInstalls(std::vector &idles, HKEY hive, REGSAM flags) +HRESULT ReadAllIdleInstalls(std::vector &idles, HKEY hive, LPCWSTR root, REGSAM flags) { HKEY hkPython = NULL; HRESULT hr = S_OK; - LSTATUS err = RegOpenKeyExW(hive, L"Software\\Python", 0, KEY_READ | flags, &hkPython); + LSTATUS err = RegOpenKeyExW(hive, root ? root : L"", 0, KEY_READ | flags, &hkPython); for (DWORD i = 0; !err && hr == S_OK && i < 64; ++i) { wchar_t name[512]; @@ -349,12 +338,12 @@ class DECLSPEC_UUID(CLSID_IDLE_COMMAND) IdleCommand iconPath += L",-4"; } - hr = ReadAllIdleInstalls(idles, HKEY_LOCAL_MACHINE, KEY_WOW64_32KEY); + hr = ReadAllIdleInstalls(idles, HKEY_LOCAL_MACHINE, L"Software\\Python", KEY_WOW64_32KEY); if (SUCCEEDED(hr)) { - hr = ReadAllIdleInstalls(idles, HKEY_LOCAL_MACHINE, KEY_WOW64_64KEY); + hr = ReadAllIdleInstalls(idles, HKEY_LOCAL_MACHINE, L"Software\\Python", KEY_WOW64_64KEY); } if (SUCCEEDED(hr)) { - hr = ReadAllIdleInstalls(idles, HKEY_CURRENT_USER, 0); + hr = ReadAllIdleInstalls(idles, HKEY_CURRENT_USER, L"Software\\Python", 0); } if (FAILED(hr)) { @@ -365,6 +354,29 @@ class DECLSPEC_UUID(CLSID_IDLE_COMMAND) IdleCommand } } + #ifdef PYSHELLEXT_TEST + IdleCommand(HKEY hive, LPCWSTR root) : title(L"Edit in &IDLE") + { + HRESULT hr; + + DWORD cch = 260; + while (iconPath.size() < cch) { + iconPath.resize(cch); + cch = GetModuleFileNameW(hModule, iconPath.data(), iconPath.size()); + } + iconPath.resize(cch); + if (cch) { + iconPath += L",-4"; + } + + hr = ReadAllIdleInstalls(idles, hive, root, 0); + + if (FAILED(hr)) { + idles.clear(); + } + } + #endif + // IExplorerCommand IFACEMETHODIMP GetTitle(IShellItemArray *psiItemArray, LPWSTR *ppszName) { @@ -427,6 +439,20 @@ class DECLSPEC_UUID(CLSID_IDLE_COMMAND) IdleCommand CoCreatableClass(IdleCommand); +#ifdef PYSHELLEXT_TEST +IExplorerCommand *MakeLaunchCommand(std::wstring title, std::wstring exe, std::wstring idle) +{ + IdleData data = { .title = title, .exe = exe, .idle = idle }; + return Make(data).Detach(); +} + + +IExplorerCommand *MakeIdleCommand(HKEY hive, LPCWSTR root) +{ + return Make(hive, root).Detach(); +} +#endif + STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, _COM_Outptr_ void** ppv) { @@ -439,7 +465,7 @@ STDAPI DllCanUnloadNow() return Module::GetModule().Terminate() ? S_OK : S_FALSE; } - +#ifndef PYSHELLEXT_TEST STDAPI_(BOOL) DllMain(_In_opt_ HINSTANCE hinst, DWORD reason, _In_opt_ void*) { if (reason == DLL_PROCESS_ATTACH) { @@ -448,3 +474,4 @@ STDAPI_(BOOL) DllMain(_In_opt_ HINSTANCE hinst, DWORD reason, _In_opt_ void*) } return TRUE; } +#endif diff --git a/src/pyshellext/shellext.h b/src/pyshellext/shellext.h new file mode 100644 index 0000000..12e1219 --- /dev/null +++ b/src/pyshellext/shellext.h @@ -0,0 +1,23 @@ +#pragma once + +#include +#include +#include +#include + +#include +#include + +LRESULT RegReadStr(HKEY key, LPCWSTR valueName, std::wstring& result); + +struct IdleData { + std::wstring title; + std::wstring exe; + std::wstring idle; +}; + +HRESULT ReadIdleInstalls(std::vector &idles, HKEY hkPython, LPCWSTR company, REGSAM flags); +HRESULT ReadAllIdleInstalls(std::vector &idles, HKEY hive, LPCWSTR root, REGSAM flags); + +IExplorerCommand *MakeIdleCommand(HKEY hive, LPCWSTR root); +IExplorerCommand *MakeLaunchCommand(std::wstring title, std::wstring exe, std::wstring idle); diff --git a/src/pyshellext/shellext_test.cpp b/src/pyshellext/shellext_test.cpp new file mode 100644 index 0000000..a574f90 --- /dev/null +++ b/src/pyshellext/shellext_test.cpp @@ -0,0 +1,241 @@ +#include + +#include "shellext.h" +#include "..\_native\helpers.h" + +extern "C" { + +PyObject *shellext_RegReadStr(PyObject *, PyObject *args, PyObject *) +{ + HKEY hkey; + PyObject *hkeyObj; + wchar_t *valueName; + if (!PyArg_ParseTuple(args, "OO&", &hkeyObj, as_utf16, &valueName)) { + return NULL; + } + if (!PyLong_AsNativeBytes(hkeyObj, &hkey, sizeof(hkey), -1)) { + PyMem_Free(valueName); + return NULL; + } + + PyObject *r = NULL; + std::wstring result; + int err = (int)RegReadStr(hkey, valueName, result); + if (err) { + PyErr_SetFromWindowsErr(err); + } else { + r = PyUnicode_FromWideChar(result.data(), result.size()); + } + + PyMem_Free(valueName); + return r; +} + + +PyObject *shellext_ReadIdleInstalls(PyObject *, PyObject *args, PyObject *) +{ + HKEY hkey; + wchar_t *company; + REGSAM flags; + PyObject *hkeyObj, *flagsObj; + if (!PyArg_ParseTuple(args, "OO&O", &hkeyObj, as_utf16, &company, &flagsObj)) { + return NULL; + } + if (!PyLong_AsNativeBytes(hkeyObj, &hkey, sizeof(hkey), -1) || + !PyLong_AsNativeBytes(flagsObj, &flags, sizeof(flags), -1)) { + PyMem_Free(company); + return NULL; + } + + PyObject *r = NULL; + std::vector result; + HRESULT hr = ReadIdleInstalls(result, hkey, company, flags); + + if (FAILED(hr)) { + PyErr_SetFromWindowsErr((int)hr); + } else { + r = PyList_New(0); + for (auto &i : result) { + PyObject *o = Py_BuildValue("uuu", i.title.c_str(), i.exe.c_str(), i.idle.c_str()); + if (!o) { + Py_CLEAR(r); + break; + } + if (PyList_Append(r, o) < 0) { + Py_DECREF(o); + Py_CLEAR(r); + break; + } + Py_DECREF(o); + } + } + + PyMem_Free(company); + return r; +} + + +PyObject *shellext_ReadAllIdleInstalls(PyObject *, PyObject *args, PyObject *) +{ + HKEY hkey; + REGSAM flags; + PyObject *hkeyObj, *flagsObj; + if (!PyArg_ParseTuple(args, "OO", &hkeyObj, &flagsObj)) { + return NULL; + } + if (!PyLong_AsNativeBytes(hkeyObj, &hkey, sizeof(hkey), -1) || + !PyLong_AsNativeBytes(flagsObj, &flags, sizeof(flags), -1)) { + return NULL; + } + + PyObject *r = NULL; + std::vector result; + HRESULT hr = ReadAllIdleInstalls(result, hkey, NULL, flags); + + if (FAILED(hr)) { + PyErr_SetFromWindowsErr((int)hr); + } else { + r = PyList_New(0); + for (auto &i : result) { + PyObject *o = Py_BuildValue("uuu", i.title.c_str(), i.exe.c_str(), i.idle.c_str()); + if (!o) { + Py_CLEAR(r); + break; + } + if (PyList_Append(r, o) < 0) { + Py_DECREF(o); + Py_CLEAR(r); + break; + } + Py_DECREF(o); + } + } + + return r; +} + + +PyObject *shellext_PassthroughTitle(PyObject *, PyObject *args, PyObject *) +{ + wchar_t *value; + if (!PyArg_ParseTuple(args, "O&", as_utf16, &value)) { + return NULL; + } + + PyObject *r = NULL; + IExplorerCommand *cmd = MakeLaunchCommand(value, L"", L""); + wchar_t *title; + HRESULT hr = cmd->GetTitle(NULL, &title); + if (SUCCEEDED(hr)) { + r = PyUnicode_FromWideChar(title, -1); + CoTaskMemFree((void*)title); + } else { + PyErr_SetFromWindowsErr((int)hr); + } + cmd->Release(); + return r; +} + + +PyObject *shellext_IdleCommand(PyObject *, PyObject *args, PyObject *) +{ + HKEY hkey; + REGSAM flags; + PyObject *hkeyObj, *flagsObj; + if (!PyArg_ParseTuple(args, "O", &hkeyObj)) { + return NULL; + } + if (!PyLong_AsNativeBytes(hkeyObj, &hkey, sizeof(hkey), -1)) { + return NULL; + } + + IExplorerCommand *cmd = MakeIdleCommand(hkey, NULL); + IEnumExplorerCommand *enm = NULL; + PyObject *r = PyList_New(0); + PyObject *o; + wchar_t *s; + HRESULT hr; + ULONG fetched; + + hr = cmd->GetTitle(NULL, &s); + if (SUCCEEDED(hr)) { + o = PyUnicode_FromWideChar(s, -1); + if (!o || PyList_Append(r, o) < 0) { + goto abort; + } + Py_CLEAR(o); + CoTaskMemFree((void *)s); + s = NULL; + } else { + goto abort; + } + + hr = cmd->GetIcon(NULL, &s); + if (SUCCEEDED(hr)) { + o = PyUnicode_FromWideChar(s, -1); + if (!o || PyList_Append(r, o) < 0) { + goto abort; + } + Py_CLEAR(o); + CoTaskMemFree((void *)s); + s = NULL; + } else { + goto abort; + } + + hr = cmd->EnumSubCommands(&enm); + cmd->Release(); + cmd = NULL; + if (FAILED(hr)) { + goto abort; + } + + while ((hr = enm->Next(1, &cmd, &fetched)) == S_OK) { + if (fetched != 1) { + PyErr_SetString(PyExc_RuntimeError, "'fetched' was not 1"); + goto abort; + } + + hr = cmd->GetTitle(NULL, &s); + if (SUCCEEDED(hr)) { + o = PyUnicode_FromWideChar(s, -1); + if (!o || PyList_Append(r, o) < 0) { + goto abort; + } + Py_CLEAR(o); + CoTaskMemFree((void *)s); + s = NULL; + } else { + goto abort; + } + + cmd->Release(); + cmd = NULL; + } + if (FAILED(hr)) { + goto abort; + } + + enm->Release(); + enm = NULL; + + return r; + +abort: + Py_XDECREF(o); + Py_XDECREF(r); + CoTaskMemFree((void *)s); + if (enm) { + enm->Release(); + } + if (cmd) { + cmd->Release(); + } + if (FAILED(hr)) { + PyErr_SetFromWindowsErr((int)hr); + } + return NULL; +} + + +} diff --git a/tests/conftest.py b/tests/conftest.py index fd844dd..b8cee5e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,6 +4,7 @@ import re import subprocess import sys +import winreg from pathlib import Path @@ -104,3 +105,48 @@ def get_installs(self): @pytest.fixture def fake_config(): return FakeConfig() + + +REG_TEST_ROOT = r"Software\Python\PyManagerTesting" + + +class RegistryFixture: + def __init__(self, hive, root): + self.hive = hive + self.root = root + self.key = None + + def __enter__(self): + self.key = winreg.CreateKey(self.hive, self.root) + return self + + def __exit__(self, *exc): + if self.key: + self.key.Close() + from manage.pep514utils import _reg_rmtree + _reg_rmtree(self.hive, self.root) + + def setup(self, _subkey=None, **keys): + if not _subkey: + _subkey = self.key + for k, v in keys.items(): + if isinstance(v, dict): + with winreg.CreateKey(_subkey, k) as subkey: + self.setup(subkey, **v) + elif isinstance(v, str): + winreg.SetValueEx(_subkey, k, None, winreg.REG_SZ, v) + elif isinstance(v, (bytes, bytearray)): + winreg.SetValueEx(_subkey, k, None, winreg.REG_BINARY, v) + elif isinstance(v, int): + if v.bit_count() < 32: + winreg.SetValueEx(_subkey, k, None, winreg.REG_DWORD, v) + else: + winreg.SetValueEx(_subkey, k, None, winreg.REG_QWORD, v) + else: + raise TypeError("unsupported type in registry") + + +@pytest.fixture(scope='function') +def registry(): + with RegistryFixture(winreg.HKEY_CURRENT_USER, REG_TEST_ROOT) as key: + yield key diff --git a/tests/test_pep514utils.py b/tests/test_pep514utils.py index 80f8bf7..994b151 100644 --- a/tests/test_pep514utils.py +++ b/tests/test_pep514utils.py @@ -3,37 +3,8 @@ from manage import pep514utils -REG_TEST_ROOT = r"Software\Python\PyManagerTesting" - -@pytest.fixture(scope='function') -def registry(): - try: - with winreg.CreateKey(winreg.HKEY_CURRENT_USER, REG_TEST_ROOT) as key: - yield key - finally: - pep514utils._reg_rmtree(winreg.HKEY_CURRENT_USER, REG_TEST_ROOT) - - -def init_reg(registry, **keys): - for k, v in keys.items(): - if isinstance(v, dict): - with winreg.CreateKey(registry, k) as subkey: - init_reg(subkey, **v) - elif isinstance(v, str): - winreg.SetValueEx(registry, k, None, winreg.REG_SZ, v) - elif isinstance(v, (bytes, bytearray)): - winreg.SetValueEx(registry, k, None, winreg.REG_BINARY, v) - elif isinstance(v, int): - if v.bit_count() < 32: - winreg.SetValueEx(registry, k, None, winreg.REG_DWORD, v) - else: - winreg.SetValueEx(registry, k, None, winreg.REG_QWORD, v) - else: - raise TypeError("unsupported type in registry") - - def test_is_tag_managed(registry, tmp_path): - init_reg(registry, Company={ + registry.setup(Company={ "1.0": {"InstallPath": {"": str(tmp_path)}}, "2.0": {"InstallPath": {"": str(tmp_path)}, "ManagedByPyManager": 0}, "2.1": {"InstallPath": {"": str(tmp_path)}, "ManagedByPyManager": 1}, @@ -42,13 +13,13 @@ def test_is_tag_managed(registry, tmp_path): "3.0.1": {"": "Also in the way here"}, }) - assert not pep514utils._is_tag_managed(registry, r"Company\1.0") - assert not pep514utils._is_tag_managed(registry, r"Company\2.0") - assert pep514utils._is_tag_managed(registry, r"Company\2.1") + assert not pep514utils._is_tag_managed(registry.key, r"Company\1.0") + assert not pep514utils._is_tag_managed(registry.key, r"Company\2.0") + assert pep514utils._is_tag_managed(registry.key, r"Company\2.1") - assert not pep514utils._is_tag_managed(registry, r"Company\3.0") + assert not pep514utils._is_tag_managed(registry.key, r"Company\3.0") with pytest.raises(FileNotFoundError): - winreg.OpenKey(registry, r"Company\3.0.2") - assert pep514utils._is_tag_managed(registry, r"Company\3.0", creating=True) - with winreg.OpenKey(registry, r"Company\3.0.2"): + winreg.OpenKey(registry.key, r"Company\3.0.2") + assert pep514utils._is_tag_managed(registry.key, r"Company\3.0", creating=True) + with winreg.OpenKey(registry.key, r"Company\3.0.2"): pass diff --git a/tests/test_shellext.py b/tests/test_shellext.py new file mode 100644 index 0000000..d865d25 --- /dev/null +++ b/tests/test_shellext.py @@ -0,0 +1,103 @@ +import pytest +import sys +import winreg + +import _shellext_test as SE + +def test_RegReadStr(): + with winreg.OpenKey(winreg.HKEY_CURRENT_USER, "Volatile Environment") as key: + assert SE.shellext_RegReadStr(key.handle, "USERPROFILE") + with pytest.raises(FileNotFoundError): + assert SE.shellext_RegReadStr(key.handle, "a made up name that hopefully doesn't exist") + with winreg.OpenKey(winreg.HKEY_CURRENT_USER, "Environment") as key: + # PATH should be REG_EXPAND_SZ, which is not supported + with pytest.raises(OSError) as ex: + SE.shellext_RegReadStr(key.handle, "PATH") + assert ex.value.winerror == 13 + + +class IdleReg: + def __init__(self, registry, tmp_path): + self.registry = registry + self.hkey = registry.key.handle + self.tmp_path = tmp_path + python_exe = tmp_path / "python.exe" + idle_pyw = tmp_path / "Lib/idlelib/idle.pyw" + self.python_exe = str(python_exe) + self.idle_pyw = str(idle_pyw) + + python_exe.parent.mkdir(parents=True, exist_ok=True) + idle_pyw.parent.mkdir(parents=True, exist_ok=True) + python_exe.write_bytes(b"") + idle_pyw.write_bytes(b"") + + registry.setup( + PythonCore={ + "1.0": { + "DisplayName": "PythonCore-1.0", + "InstallPath": { + "": str(tmp_path), + } + }, + }, + # Even if all the pieces are there, we won't pick up non-PythonCore + # unless they specify IdlePath + NotPythonCore={ + "1.0": { + "DisplayName": "NotPythonCore-1.0", + "InstallPath": { + "": str(tmp_path), + } + }, + "2.0": { + "DisplayName": "NotPythonCore-2.0", + "InstallPath": { + "": str(tmp_path), + "WindowedExecutablePath": str(python_exe), + "IdlePath": str(idle_pyw), + } + }, + }, + ) + + self.pythoncore_1_0 = ("PythonCore-1.0", self.python_exe, self.idle_pyw) + self.pythoncore = [self.pythoncore_1_0] + # NotPythonCore-1.0 should never get returned + self.notpythoncore_1_0 = ("NotPythonCore-1.0", self.python_exe, self.idle_pyw) + self.notpythoncore_2_0 = ("NotPythonCore-2.0", self.python_exe, self.idle_pyw) + self.notpythoncore = [self.notpythoncore_2_0] + self.all = [*self.notpythoncore, *self.pythoncore] + + +@pytest.fixture(scope='function') +def idle_reg(registry, tmp_path): + return IdleReg(registry, tmp_path) + + +def test_ReadIdleInstalls(idle_reg): + inst = SE.shellext_ReadIdleInstalls(idle_reg.hkey, "PythonCore", 0) + assert inst == idle_reg.pythoncore + inst = SE.shellext_ReadIdleInstalls(idle_reg.hkey, "NotPythonCore", 0) + assert inst == idle_reg.notpythoncore + + +def test_ReadAllIdleInstalls(idle_reg): + inst = SE.shellext_ReadAllIdleInstalls(idle_reg.hkey, 0) + assert inst == [ + *idle_reg.notpythoncore, + *idle_reg.pythoncore, + ] + + +def test_PassthroughTitle(): + assert "Test" == SE.shellext_PassthroughTitle("Test") + assert "Test \u0ABC" == SE.shellext_PassthroughTitle("Test \u0ABC") + + +def test_IdleCommand(idle_reg): + data = SE.shellext_IdleCommand(idle_reg.hkey) + assert data == [ + "Edit in &IDLE", + f"{sys._base_executable},-4", + *(i[0] for i in reversed(idle_reg.all)), + ]