diff --git a/Pipfile b/Pipfile index 8dfb404..720f73d 100644 --- a/Pipfile +++ b/Pipfile @@ -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] diff --git a/pew/_utils.py b/pew/_utils.py index d1d8e0a..e345937 100644 --- a/pew/_utils.py +++ b/pew/_utils.py @@ -30,7 +30,12 @@ def to_unicode(x): return x.decode(encoding) else: NamedTemporaryFile = _ntf - to_unicode = str + + def to_unicode(x): + if isinstance(x, (bytes, bytearray)): + x = x.decode(encoding) + return str(x) + def check_path(): parent = os.path.dirname diff --git a/pew/_win_utils.py b/pew/_win_utils.py new file mode 100644 index 0000000..ac37763 --- /dev/null +++ b/pew/_win_utils.py @@ -0,0 +1,140 @@ +# -*- 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 +from ._utils import to_unicode + +ERROR_NO_MORE_FILES = 18 +INVALID_HANDLE_VALUE = c_void_p(-1).value +SHELL_NAMES = ['cmd', 'powershell', 'pwsh', 'cmder'] +global SHELL_NAMES + +if os.environ.get('RUNNING_CI'): + SHELL_NAMES.append('python') + + +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': to_unicode(pe.szExeFile) + } + if pe.th32ParentProcessID: + pids[pe.th32ProcessID]['parent_pid'] = pe.th32ParentProcessID + pe = Process32Next(h_process, pe) + + return pids + + +def _get_executable(process_dict): + if hasattr(process_dict, 'keys'): + executable = process_dict.get('executable') + if executable: + return executable.lower().rsplit('.', 1)[0] + return '' + + +def get_shell(pid=None, max_depth=6): + """Get the shell that the supplied pid or os.getpid() is running in. + """ + if not pid: + pid = os.getpid() + processes = get_all_processes() + own_exe = to_unicode(processes[pid]['executable']) + + def check_parent(pid, lvl=0, processes=None): + if not processes: + processes = get_all_processes() + if pid not in processes: + # see if the process graph has updated since the last check + processes = get_all_processes() + if pid not in processes: + return + ppid = processes[pid].get('parent_pid') + if ppid and _get_executable(processes.get(ppid)) in SHELL_NAMES: + return to_unicode(processes[ppid]['executable']) + if lvl >= max_depth: + return + return check_parent(ppid, lvl=lvl+1, processes=processes) + + if os.environ.get('RUNNING_CI'): + global SHELL_NAMES + max_depth = 4 + if 'python' not in SHELL_NAMES: + SHELL_NAMES = ['python'] + SHELL_NAMES + + if _get_executable(processes.get(pid)) in SHELL_NAMES: + return own_exe + return check_parent(pid, processes=processes) or own_exe diff --git a/pew/pew.py b/pew/pew.py index cc5a042..c9db11e 100644 --- a/pew/pew.py +++ b/pew/pew.py @@ -36,7 +36,7 @@ InstallCommand = ListPythons = LocatePython = UninstallCommand = \ lambda : sys.exit('Command not supported on this platform') - import psutil + from ._win_utils import get_shell from pew._utils import (check_call, invoke, expandpath, own, env_bin_dir, check_path, temp_environ, NamedTemporaryFile, to_unicode) @@ -52,8 +52,9 @@ else: default_home = os.path.join( os.environ.get('XDG_DATA_HOME', '~/.local/share'), 'virtualenvs') -workon_home = expandpath( - os.environ.get('WORKON_HOME', default_home)) + +def get_workon_home(): + return expandpath(os.environ.get('WORKON_HOME', default_home)) def makedirs_and_symlink_if_needed(workon_home): @@ -96,7 +97,7 @@ def deploy_completions(): def get_project_dir(env): - project_file = workon_home / env / '.project' + project_file = get_workon_home() / env / '.project' if project_file.exists(): with project_file.open() as f: project_dir = f.readline().strip() @@ -113,7 +114,7 @@ def unsetenv(key): def compute_path(env): - envdir = workon_home / env + envdir = get_workon_home() / env return os.pathsep.join([ str(envdir / env_bin_dir), os.environ['PATH'], @@ -127,7 +128,7 @@ def inve(env, command, *args, **kwargs): # we don't strictly need to restore the environment, since pew runs in # its own process, but it feels like the right thing to do with temp_environ(): - os.environ['VIRTUAL_ENV'] = str(workon_home / env) + os.environ['VIRTUAL_ENV'] = str(get_workon_home() / env) os.environ['PATH'] = compute_path(env) unsetenv('PYTHONHOME') @@ -151,7 +152,16 @@ def fork_shell(env, shellcmd, cwd): if 'VIRTUAL_ENV' in os.environ: err("Be aware that this environment will be nested on top " "of '%s'" % Path(os.environ['VIRTUAL_ENV']).name) - inve(env, *shellcmd, cwd=cwd) + try: + inve(env, *shellcmd, cwd=cwd) + except CalledProcessError: + # These shells report errors when the last command executed in the + # subshell in an error. This causes the subprocess to fail, which is + # not what we want. Stay silent for them, there's nothing we can do. + shell_name, _ = os.path.splitext(os.path.basename(shellcmd[0])) + suppress_error = shell_name.lower() in ('cmd', 'powershell', 'pwsh') + if not suppress_error: + raise def fork_bash(env, cwd): @@ -184,7 +194,7 @@ def _detect_shell(): if 'CMDER_ROOT' in os.environ: shell = 'Cmder' elif windows: - shell = psutil.Process(os.getpid()).parent().parent().name() + shell = get_shell(os.getpid()) else: shell = 'sh' return shell @@ -193,7 +203,7 @@ def shell(env, cwd=None): env = str(env) shell = _detect_shell() shell_name = Path(shell).stem - if shell_name not in ('Cmder', 'bash', 'elvish', 'powershell', 'klingon', 'cmd'): + if shell_name not in ('Cmder', 'bash', 'elvish', 'powershell', 'pwsh', 'klingon', 'cmd'): # On Windows the PATH is usually set with System Utility # so we won't worry about trying to check mistakes there shell_check = (sys.executable + ' -c "from pew.pew import ' @@ -216,7 +226,7 @@ def mkvirtualenv(envname, python=None, packages=[], project=None, if python: rest = ["--python=%s" % python] + rest - path = (workon_home / envname).absolute() + path = (get_workon_home() / envname).absolute() try: check_call([sys.executable, "-m", "virtualenv", str(path)] + rest) @@ -265,7 +275,7 @@ def new_cmd(argv): def rmvirtualenvs(envs): error_happened = False for env in envs: - env = workon_home / env + env = get_workon_home() / env if os.environ.get('VIRTUAL_ENV') == str(env): err("ERROR: You cannot remove the active environment (%s)." % env) error_happened = True @@ -295,7 +305,7 @@ def packages(site_packages): def showvirtualenv(env): columns, _ = get_terminal_size() pkgs = sorted(packages(sitepackages_dir(env))) - env_python = workon_home / env / env_bin_dir / 'python' + env_python = get_workon_home() / env / env_bin_dir / 'python' l = len(env) + 2 version = invoke(str(env_python), '-V') version = ' - '.join((version.out + version.err).splitlines()) @@ -317,8 +327,8 @@ def show_cmd(argv): def lsenvs(): - return sorted(set(env.parts[-3] for env in - workon_home.glob(os.path.join('*', env_bin_dir, 'python*')))) + items = get_workon_home().glob(os.path.join('*', env_bin_dir, 'python*')) + return sorted(set(env.parts[-3] for env in items)) def lsvirtualenv(verbose): @@ -347,7 +357,7 @@ def parse_envname(argv, no_arg_callback): env = argv[0] if env.startswith('/'): sys.exit("ERROR: Invalid environment name '{0}'.".format(env)) - if not (workon_home / env).exists(): + if not (get_workon_home() / env).exists(): sys.exit("ERROR: Environment '{0}' does not exist. Create it with \ 'pew new {0}'.".format(env)) else: @@ -375,7 +385,7 @@ def sitepackages_dir(env=os.environ.get('VIRTUAL_ENV')): if not env: sys.exit('ERROR: no virtualenv active') else: - env_python = workon_home / env / env_bin_dir / 'python' + env_python = get_workon_home() / env / env_bin_dir / 'python' return Path(invoke(str(env_python), '-c', 'import distutils; \ print(distutils.sysconfig.get_python_lib())').out) @@ -462,6 +472,7 @@ def cp_cmd(argv): def copy_virtualenv_project(source, target): source = expandpath(source) + workon_home = get_workon_home() if not source.exists(): source = workon_home / source if not source.exists(): @@ -493,7 +504,7 @@ def rename_cmd(argv): def setvirtualenvproject(env, project): print('Setting project for {0} to {1}'.format(env, project)) - with (workon_home / env / '.project').open('wb') as prj: + with (get_workon_home() / env / '.project').open('wb') as prj: prj.write(str(project).encode()) @@ -505,7 +516,7 @@ def setproject_cmd(argv): env = args.get(0, os.environ.get('VIRTUAL_ENV')) if not env: sys.exit('pew setproject [virtualenv] [project_path]') - if not (workon_home / env).exists(): + if not (get_workon_home() / env).exists(): sys.exit("Environment '%s' doesn't exist." % env) if not os.path.isdir(project): sys.exit('pew setproject: %s does not exist' % project) @@ -515,7 +526,7 @@ def setproject_cmd(argv): def mkproject_cmd(argv): """Create a new project directory and its associated virtualenv.""" if '-l' in argv or '--list' in argv: - templates = [t.name[9:] for t in workon_home.glob("template_*")] + templates = [t.name[9:] for t in get_workon_home().glob("template_*")] print("Available project templates:", *templates, sep='\n') return @@ -545,7 +556,7 @@ def mkproject_cmd(argv): project.mkdir() for template_name in args.templates: - template = workon_home / ("template_" + template_name) + template = get_workon_home() / ("template_" + template_name) inve(args.envname, str(template), args.envname, str(project)) if args.activate: shell(args.envname, cwd=str(project)) @@ -555,7 +566,7 @@ def mktmpenv_cmd(argv): """Create a temporary virtualenv.""" parser = mkvirtualenv_argparser() env = '.' - while (workon_home / env).exists(): + while (get_workon_home() / env).exists(): env = hex(random.getrandbits(64))[2:-1] args, rest = parser.parse_known_args(argv) @@ -577,10 +588,10 @@ def wipeenv_cmd(argv): if not env: sys.exit('ERROR: no virtualenv active') - elif not (workon_home / env).exists(): + elif not (get_workon_home() / env).exists(): sys.exit("ERROR: Environment '{0}' does not exist.".format(env)) else: - env_pip = str(workon_home / env / env_bin_dir / 'pip') + env_pip = str(get_workon_home() / env / env_bin_dir / 'pip') all_pkgs = set(invoke(env_pip, 'freeze').out.splitlines()) pkgs = set(p for p in all_pkgs if len(p.split("==")) == 2) ignored = sorted(all_pkgs - pkgs) @@ -626,7 +637,7 @@ def restore_cmd(argv): sys.exit('You must provide a valid virtualenv to target') env = argv[0] - path = workon_home / env + path = get_workon_home() / env py = path / env_bin_dir / ('python.exe' if windows else 'python') exact_py = py.resolve().name @@ -636,7 +647,7 @@ def restore_cmd(argv): def dir_cmd(argv): """Print the path for the virtualenv directory""" env = parse_envname(argv, lambda : sys.exit('You must provide a valid virtualenv to target')) - print(workon_home / env) + print(get_workon_home() / env) def install_cmd(argv): @@ -748,7 +759,7 @@ def print_commands(cmds): def pew(): - first_run = makedirs_and_symlink_if_needed(workon_home) + first_run = makedirs_and_symlink_if_needed(get_workon_home()) if first_run and sys.stdin.isatty(): first_run_setup() diff --git a/requirements.txt b/requirements.txt index 565f66d..f17ca6f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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' \ No newline at end of file +stdeb ; sys_platform == 'linux' diff --git a/setup.py b/setup.py index dc2e0de..cd113e1 100644 --- a/setup.py +++ b/setup.py @@ -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' ] diff --git a/tests/test_workon.py b/tests/test_workon.py index 488b176..1f05e99 100644 --- a/tests/test_workon.py +++ b/tests/test_workon.py @@ -5,9 +5,9 @@ from pew.pew import _detect_shell from pew._utils import temp_environ, invoke_pew as invoke from utils import skip_windows - import pytest + check_env = [sys.executable, '-c', "import os; print(os.environ['VIRTUAL_ENV'])"] check_cwd = [sys.executable, '-c', "import pathlib; print(pathlib.Path().absolute())"] @@ -18,7 +18,11 @@ def test_detect_shell(): except KeyError: pass if sys.platform == 'win32': - assert _detect_shell() in ['python.exe'] + from pew._win_utils import SHELL_NAMES as WIN_SHELL_NAMES + WIN_SHELL_NAMES.append('python') + win_shells = WIN_SHELL_NAMES + ['{}.exe'.format(s) for s in WIN_SHELL_NAMES] + os.environ['RUNNING_CI'] = '1' + assert _detect_shell() in win_shells else: assert _detect_shell() == 'sh' os.environ['SHELL'] = 'foo'