Skip to content

Commit

Permalink
feat: added helper method to receive win folders
Browse files Browse the repository at this point in the history
  • Loading branch information
finswimmer committed Jun 30, 2022
1 parent ceb3586 commit 0f19fe7
Showing 1 changed file with 69 additions and 0 deletions.
69 changes: 69 additions & 0 deletions src/poetry/utils/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import re
import shutil
import stat
import sys
import tempfile

from contextlib import contextmanager
Expand All @@ -12,6 +13,7 @@
from typing import Any
from typing import Iterator
from typing import Mapping
from typing import cast

from poetry.utils.constants import REQUESTS_TIMEOUT

Expand Down Expand Up @@ -171,3 +173,70 @@ def safe_extra(extra: str) -> str:
https://github.com/pypa/setuptools/blob/452e13c/pkg_resources/__init__.py#L1423-L1431.
"""
return re.sub("[^A-Za-z0-9.-]+", "_", extra).lower()


def _get_win_folder_from_registry(csidl_name: str) -> str:
if sys.platform != "win32":
raise RuntimeError("Method can only be called on Windows.")

import winreg as _winreg

shell_folder_name = {
"CSIDL_APPDATA": "AppData",
"CSIDL_COMMON_APPDATA": "Common AppData",
"CSIDL_LOCAL_APPDATA": "Local AppData",
"CSIDL_PROGRAM_FILES": "Program Files",
}[csidl_name]

key = _winreg.OpenKey(
_winreg.HKEY_CURRENT_USER,
r"Software\Microsoft\Windows\CurrentVersion\Explorer\Shell Folders",
)
dir, type = _winreg.QueryValueEx(key, shell_folder_name)

return cast(str, dir)


def _get_win_folder_with_ctypes(csidl_name: str) -> str:
if sys.platform != "win32":
raise RuntimeError("Method can only be called on Windows.")

import ctypes

csidl_const = {
"CSIDL_APPDATA": 26,
"CSIDL_COMMON_APPDATA": 35,
"CSIDL_LOCAL_APPDATA": 28,
"CSIDL_PROGRAM_FILES": 38,
}[csidl_name]

buf = ctypes.create_unicode_buffer(1024)
ctypes.windll.shell32.SHGetFolderPathW(None, csidl_const, None, 0, buf)

# Downgrade to short path name if have highbit chars. See
# <http://bugs.activestate.com/show_bug.cgi?id=85099>.
has_high_char = False
for c in buf:
if ord(c) > 255:
has_high_char = True
break
if has_high_char:
buf2 = ctypes.create_unicode_buffer(1024)
if ctypes.windll.kernel32.GetShortPathNameW(buf.value, buf2, 1024):
buf = buf2

return buf.value


def get_win_folder(csidl_name: str) -> Path:
if sys.platform == "win32":
try:
from ctypes import windll # noqa: F401

_get_win_folder = _get_win_folder_with_ctypes
except ImportError:
_get_win_folder = _get_win_folder_from_registry

return Path(_get_win_folder(csidl_name))

raise RuntimeError("Method can only be called on Windows.")

0 comments on commit 0f19fe7

Please sign in to comment.