From 048ef3b3174350a302b6d7ebccb8480432a1b7cb Mon Sep 17 00:00:00 2001
From: Amos Bird <amosbird@gmail.com>
Date: Tue, 14 May 2024 22:23:36 +0800
Subject: [PATCH] Support --env argument for @ run

---
 kitty/launch.py | 15 ++++++++++-----
 kitty/rc/run.py | 20 ++++++++++++++++----
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/kitty/launch.py b/kitty/launch.py
index e14029588..00eb5ccd4 100644
--- a/kitty/launch.py
+++ b/kitty/launch.py
@@ -30,6 +30,15 @@ class LaunchSpec(NamedTuple):
     args: List[str]
 
 
+env_docs = '''\
+type=list
+Environment variables to set in the child process. Can be specified multiple
+times to set different environment variables. Syntax: :code:`name=value`. Using
+:code:`name=` will set to empty string and just :code:`name` will remove the
+environment variable.
+'''
+
+
 remote_control_password_docs = '''\
 type=list
 Restrict the actions remote control is allowed to take. This works like
@@ -124,11 +133,7 @@ refers to the process that was originally started when the window was created.
 
 
 --env
-type=list
-Environment variables to set in the child process. Can be specified multiple
-times to set different environment variables. Syntax: :code:`name=value`. Using
-:code:`name=` will set to empty string and just :code:`name` will remove the
-environment variable.
+{env_docs}
 
 
 --var
diff --git a/kitty/rc/run.py b/kitty/rc/run.py
index 88540c4be..e28beb945 100644
--- a/kitty/rc/run.py
+++ b/kitty/rc/run.py
@@ -3,10 +3,11 @@
 
 import sys
 from base64 import standard_b64decode, standard_b64encode
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, Dict
 
-from kitty.launch import remote_control_password_docs
+from kitty.launch import env_docs, remote_control_password_docs
 from kitty.types import AsyncResponse
+from kitty.options.utils import env as parse_env
 
 from .base import (
     ArgsType,
@@ -29,6 +30,7 @@ class Run(RemoteCommand):
     protocol_spec = __doc__ = '''
     data+/str: Chunk of STDIN data, base64 encoded no more than 4096 bytes. Must send an empty chunk to indicate end of data.
     cmdline+/list.str: The command line to run
+    env/list.str: List of environment variables of the form NAME=VALUE
     allow_remote_control/bool: A boolean indicating whether to allow remote control
     remote_control_password/list.str: A list of remote control passwords
     '''
@@ -42,6 +44,10 @@ class Run(RemoteCommand):
     )
 
     options_spec = f'''\n
+--env
+{env_docs}
+
+
 --allow-remote-control
 type=bool-set
 The executed program will have privileges to run remote control commands in kitty.
@@ -64,6 +70,7 @@ The executed program will have privileges to run remote control commands in kitt
         ret = {
             'stream_id': secrets.token_urlsafe(),
             'cmdline': args,
+            'env': opts.env,
             'allow_remote_control': opts.allow_remote_control,
             'remote_control_password': opts.remote_control_password,
             'data': '',
@@ -113,9 +120,14 @@ The executed program will have privileges to run remote control commands in kitt
                         'exit_code': exit_code, 'exit_status': exit_status,
                     })
 
+        env: Dict[str, str] = {}
+        for x in payload_get('env'):
+            for k, v in parse_env(x, env):
+                env[k] = v
+
         boss.run_background_process(
-            cmdline, stdin=stdin_data, stdout=stdout.fileno(), stderr=stderr.fileno(), notify_on_death=on_death,
-            remote_control_passwords=rcp, allow_remote_control=allow_remote_control
+            cmdline, env=env, stdin=stdin_data, stdout=stdout.fileno(), stderr=stderr.fileno(),
+            notify_on_death=on_death, remote_control_passwords=rcp, allow_remote_control=allow_remote_control
         )
         return AsyncResponse()