Skip to content

Commit

Permalink
Remove dependence on PSUtil
Browse files Browse the repository at this point in the history
- Use native ctypes for parent process traversal
- Addresses issues raised in pypa/pipenv#1587
  and microsoft/vscode-python#978
- Improves speed on windows
- Allows pipenv to remove vendored psutil which sometimes fails to find
linked python dlls
  • Loading branch information
techalchemy committed Apr 20, 2018
1 parent 345617f commit 818ab1e
Show file tree
Hide file tree
Showing 5 changed files with 110 additions and 8 deletions.
1 change: 0 additions & 1 deletion Pipfile
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ name = "pypi"
virtualenv = "==1.11.6"
"virtualenv-clone" = "==0.2.5"
"pythonz-bd" = {"version" = "*", "sys_platform" = "!='win32'" }
"psutil" = {"version" = "*", "sys_platform" = "=='win32'" }

[dev-packages]

Expand Down
107 changes: 107 additions & 0 deletions pew/_win_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
# -*- coding=utf-8 -*-
# psutil is painfully slow in win32. So to avoid adding big
# dependencies like pywin32 a ctypes based solution is preferred

# Code based on the winappdbg project http://winappdbg.sourceforge.net/
# (BSD License) - adapted from Celery
# https://github.com/celery/celery/blob/2.5-archived/celery/concurrency/processes/_win.py
import os
from ctypes import (
byref, sizeof, windll, Structure, WinError, POINTER,
c_size_t, c_char, c_void_p
)
from ctypes.wintypes import DWORD, LONG

ERROR_NO_MORE_FILES = 18
INVALID_HANDLE_VALUE = c_void_p(-1).value


class PROCESSENTRY32(Structure):
_fields_ = [
('dwSize', DWORD),
('cntUsage', DWORD),
('th32ProcessID', DWORD),
('th32DefaultHeapID', c_size_t),
('th32ModuleID', DWORD),
('cntThreads', DWORD),
('th32ParentProcessID', DWORD),
('pcPriClassBase', LONG),
('dwFlags', DWORD),
('szExeFile', c_char * 260),
]


LPPROCESSENTRY32 = POINTER(PROCESSENTRY32)


def CreateToolhelp32Snapshot(dwFlags=2, th32ProcessID=0):
hSnapshot = windll.kernel32.CreateToolhelp32Snapshot(
dwFlags,
th32ProcessID
)
if hSnapshot == INVALID_HANDLE_VALUE:
raise WinError()
return hSnapshot


def Process32First(hSnapshot):
pe = PROCESSENTRY32()
pe.dwSize = sizeof(PROCESSENTRY32)
success = windll.kernel32.Process32First(hSnapshot, byref(pe))
if not success:
if windll.kernel32.GetLastError() == ERROR_NO_MORE_FILES:
return
raise WinError()
return pe


def Process32Next(hSnapshot, pe=None):
if pe is None:
pe = PROCESSENTRY32()
pe.dwSize = sizeof(PROCESSENTRY32)
success = windll.kernel32.Process32Next(hSnapshot, byref(pe))
if not success:
if windll.kernel32.GetLastError() == ERROR_NO_MORE_FILES:
return
raise WinError()
return pe


def get_all_processes():
"""Return a dictionary of properties about all processes.
>>> get_all_processes()
{
1509: {
'parent_pid': 1201,
'executable': 'C:\\Program\\\\ Files\\Python36\\python.exe'
}
}
"""
h_process = CreateToolhelp32Snapshot()
pids = {}
pe = Process32First(h_process)
while pe:
pids[pe.th32ProcessID] = {
'executable': str(pe.szExeFile.decode('utf-8')),
}
if pe.th32ParentProcessID:
pids[pe.th32ProcessID]['parent_pid'] = pe.th32ParentProcessID
pe = Process32Next(h_process, pe)

return pids


def get_grandparent_process(pid=None):
"""Get grandparent process name of the supplied pid or os.getpid().
:param int pid: The pid to track.
:return: Name of the grandparent process.
"""
if not pid:
pid = os.getpid()
processes = get_all_processes()
ppid = processes[pid]['parent_pid']
parent = processes[ppid]
grandparent = processes[parent['parent_pid']]
return grandparent['executable']
4 changes: 2 additions & 2 deletions pew/pew.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
InstallCommand = ListPythons = LocatePython = UninstallCommand = \
lambda : sys.exit('Command not supported on this platform')

import psutil
from ._win_utils import get_grandparent_process

from pew._utils import (check_call, invoke, expandpath, own, env_bin_dir,
check_path, temp_environ, NamedTemporaryFile, to_unicode)
Expand Down Expand Up @@ -184,7 +184,7 @@ def _detect_shell():
if 'CMDER_ROOT' in os.environ:
shell = 'Cmder'
elif windows:
shell = psutil.Process(os.getpid()).parent().parent().name()
shell = get_grandparent_process(os.getpid())
else:
shell = 'sh'
return shell
Expand Down
3 changes: 1 addition & 2 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -3,5 +3,4 @@ virtualenv==1.11.6
virtualenv-clone==0.2.5
pytest==2.6.2
pythonz-bd==1.11.2 ; sys_platform != 'win32'
psutil==5.3.1 ; sys_platform == 'win32'
stdeb ; sys_platform == 'linux'
stdeb ; sys_platform == 'linux'
3 changes: 0 additions & 3 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,9 +77,6 @@ def run(self):
':python_version=="3.3"': [
'pathlib'
],
':sys_platform=="win32"': [
'psutil==5.3.1'
],
'pythonz': [
'pythonz-bd>=1.10.2'
]
Expand Down

0 comments on commit 818ab1e

Please sign in to comment.