some PEP8 and code deduplication

Author: dasmy

gpt_code_ui/kernel_program/config.py

@@ -15,4 +15,4 @@ def get_logger():
     logger = logging.getLogger(__name__)
     if "DEBUG" in os.environ:
         logger.setLevel(logging.DEBUG)
-    return logger
+    return logger

gpt_code_ui/kernel_program/kernel_manager.py

@@ -57,7 +57,7 @@ def cleanup_spawned_processes():
                         os.kill(pid, signal.CTRL_BREAK_EVENT)
                     else:
                         os.kill(pid, signal.SIGKILL)
-                    
+
                     # After successful kill, cleanup pid file
                     os.remove(fp)
 
@@ -217,11 +217,8 @@ def start_kernel(id: str):
         cwd=kernel_dir,
         env=kernel_env,
     )
-    # Write PID for caller to kill
-    str_kernel_pid = str(kernel_process.pid)
-    os.makedirs(config.KERNEL_PID_DIR, exist_ok=True)
-    with open(os.path.join(config.KERNEL_PID_DIR, str_kernel_pid + ".pid"), "w") as p:
-        p.write("kernel")
+
+    utils.store_pid(kernel_process.pid, "kernel")
 
     # Wait for kernel connection file to be written
     while True:

gpt_code_ui/kernel_program/launch_kernel.py

@@ -1,4 +1,4 @@
 if __name__ == "__main__":
     from ipykernel import kernelapp as app
 
-    app.launch_new_instance()
+    app.launch_new_instance()

gpt_code_ui/kernel_program/main.py

@@ -7,7 +7,6 @@
 import time
 
 import asyncio
-import json
 import threading
 
 from queue import Queue
@@ -60,14 +59,13 @@ def start_kernel_manager():
         'workspace',  # This will be used as part of the folder name for the workspace and to create the venv inside. Can be anything, but using 'workspace' makes file up-/download very simple
     ])
 
-    # Write PID as <pid>.pid to config.KERNEL_PID_DIR
-    os.makedirs(config.KERNEL_PID_DIR, exist_ok=True)
-    with open(os.path.join(config.KERNEL_PID_DIR, "%d.pid" % kernel_manager_process.pid), "w") as p:
-        p.write("kernel_manager")
+    utils.store_pid(kernel_manager_process.pid, "kernel_manager")
+
 
 def cleanup_kernel_program():
     kernel_manager.cleanup_spawned_processes()
 
+
 async def start_snakemq():
     global messaging
 
@@ -80,7 +78,7 @@ def on_recv(conn, ident, message):
             if message["value"] == "ready":
                 logger.debug("Kernel is ready.")
                 result_queue.put({
-                    "value":"Kernel is ready.",
+                    "value": "Kernel is ready.",
                     "type": "message"
                 })
 
@@ -100,8 +98,9 @@ def send_queued_messages():
         while True:
             if send_queue.qsize() > 0:
                 message = send_queue.get()
-                utils.send_json(messaging, 
-                    {"type": "execute", "value": message["command"]}, 
+                utils.send_json(
+                    messaging,
+                    {"type": "execute", "value": message["command"]},
                     config.IDENT_KERNEL_MANAGER
                 )
             time.sleep(0.1)
@@ -120,7 +119,7 @@ async def async_link_loop():
 
 @app.route("/api", methods=["POST", "GET"])
 def handle_request():
-   
+
     if request.method == "GET":
         # Handle GET requests by sending everything that's in the receive_queue
         results = [result_queue.get() for _ in range(result_queue.qsize())]
@@ -131,7 +130,8 @@ def handle_request():
         send_queue.put(data)
 
         return jsonify({"result": "success"})
-    
+
+
 @app.route("/restart", methods=["POST"])
 def handle_restart():
 
@@ -155,9 +155,6 @@ async def main():
 def run_flask_app():
     app.run(host="0.0.0.0", port=APP_PORT)
 
+
 if __name__ == "__main__":
     asyncio.run(main())
-
-
-
-    

gpt_code_ui/kernel_program/utils.py

@@ -1,3 +1,4 @@
+import os
 import re
 import json
 import snakemq.link
@@ -7,6 +8,7 @@
 
 import gpt_code_ui.kernel_program.config as config
 
+
 def escape_ansi(line):
     ansi_escape = re.compile(r"(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]")
     return ansi_escape.sub("", line)
@@ -16,6 +18,7 @@ def send_json(messaging, message, identity):
     message = snakemq.message.Message(json.dumps(message).encode("utf-8"), ttl=600)
     messaging.send_message(identity, message)
 
+
 def init_snakemq(ident, init_type="listen"):
     link = snakemq.link.Link()
     packeter = snakemq.packeter.Packeter(link)
@@ -26,4 +29,13 @@ def init_snakemq(ident, init_type="listen"):
         link.add_connector(("localhost", config.SNAKEMQ_PORT))
     else:
         raise Exception("Unsupported init type.")
-    return messaging, link
+    return messaging, link
+
+
+def store_pid(pid: int, process_name: str):
+    '''
+    Write PID as <pid>.pid to config.KERNEL_PID_DIR
+    '''
+    os.makedirs(config.KERNEL_PID_DIR, exist_ok=True)
+    with open(os.path.join(config.KERNEL_PID_DIR, f"{pid}.pid"), "w") as p:
+        p.write(process_name)