[compat] Add compat_setenv

[youtube-dl] / youtube_dl / compat.py

diff --git a/youtube_dl/compat.py b/youtube_dl/compat.py
index 0b6c5ca7a8ba5eb6cb064916d56a5ca8eae32003..12b53cdc88f8de7e0d37a46cc9804a193463db1e 100644 (file)
--- a/youtube_dl/compat.py
+++ b/youtube_dl/compat.py
@@ -373,6 +373,9 @@ compat_os_name = os._name if os.name == 'java' else os.name
 if sys.version_info >= (3, 0):
     compat_getenv = os.getenv
     compat_expanduser = os.path.expanduser
+
+    def compat_setenv(key, value, env=os.environ):
+        env[key] = value
 else:
     # Environment variables should be decoded with filesystem encoding.
     # Otherwise it will fail if any non-ASCII characters present (see #3854 #3217 #2918)
@@ -384,6 +387,12 @@ else:
             env = env.decode(get_filesystem_encoding())
         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
     # environment variables with filesystem encoding. We will work around this by
     # providing adjusted implementations.
@@ -604,6 +613,7 @@ __all__ = [
     'compat_os_name',
     'compat_parse_qs',
     'compat_print',
+    'compat_setenv',
     'compat_shlex_split',
     'compat_socket_create_connection',
     'compat_str',