[compat] Add compat_setenv

This commit is contained in:
Sergey M․ 2016-05-09 21:55:03 +06:00
parent 6104cc2985
commit fe40f9eef2
No known key found for this signature in database
GPG Key ID: 2C393E0F18A9236D
2 changed files with 18 additions and 0 deletions

View File

@ -13,6 +13,7 @@
from youtube_dl.utils import get_filesystem_encoding from youtube_dl.utils import get_filesystem_encoding
from youtube_dl.compat import ( from youtube_dl.compat import (
compat_getenv, compat_getenv,
compat_setenv,
compat_etree_fromstring, compat_etree_fromstring,
compat_expanduser, compat_expanduser,
compat_shlex_split, compat_shlex_split,
@ -31,6 +32,13 @@ def test_compat_getenv(self):
else test_str.encode(get_filesystem_encoding())) else test_str.encode(get_filesystem_encoding()))
self.assertEqual(compat_getenv('YOUTUBE-DL-TEST'), test_str) self.assertEqual(compat_getenv('YOUTUBE-DL-TEST'), test_str)
def test_compat_setenv(self):
test_var = 'YOUTUBE-DL-TEST'
test_str = 'тест'
compat_setenv(test_var, test_str)
compat_getenv(test_var)
self.assertEqual(compat_getenv(test_var), test_str)
def test_compat_expanduser(self): def test_compat_expanduser(self):
old_home = os.environ.get('HOME') old_home = os.environ.get('HOME')
test_str = 'C:\Documents and Settings\тест\Application Data' test_str = 'C:\Documents and Settings\тест\Application Data'

View File

@ -373,6 +373,9 @@ def compat_ord(c):
if sys.version_info >= (3, 0): if sys.version_info >= (3, 0):
compat_getenv = os.getenv compat_getenv = os.getenv
compat_expanduser = os.path.expanduser compat_expanduser = os.path.expanduser
def compat_setenv(key, value, env=os.environ):
env[key] = value
else: else:
# Environment variables should be decoded with filesystem encoding. # Environment variables should be decoded with filesystem encoding.
# Otherwise it will fail if any non-ASCII characters present (see #3854 #3217 #2918) # Otherwise it will fail if any non-ASCII characters present (see #3854 #3217 #2918)
@ -384,6 +387,12 @@ def compat_getenv(key, default=None):
env = env.decode(get_filesystem_encoding()) env = env.decode(get_filesystem_encoding())
return env return env
def compat_setenv(key, value, env=os.environ):
def encode(v):
from .utils import get_filesystem_encoding
return v.encode(get_filesystem_encoding()) if isinstance(v, compat_str) else v
env[encode(key)] = encode(value)
# HACK: The default implementations of os.path.expanduser from cpython do not decode # HACK: The default implementations of os.path.expanduser from cpython do not decode
# environment variables with filesystem encoding. We will work around this by # environment variables with filesystem encoding. We will work around this by
# providing adjusted implementations. # providing adjusted implementations.
@ -604,6 +613,7 @@ def compat_itertools_count(start=0, step=1):
'compat_os_name', 'compat_os_name',
'compat_parse_qs', 'compat_parse_qs',
'compat_print', 'compat_print',
'compat_setenv',
'compat_shlex_split', 'compat_shlex_split',
'compat_socket_create_connection', 'compat_socket_create_connection',
'compat_str', 'compat_str',