add rich based console to the file (Run ID: tripathi456_gvisor-based-python-repl_issue_9_d918a6d9)

test_tcp.py

@@ -2,6 +2,14 @@
 import socket
 import json
 import sys
+from rich.console import Console
+from rich.panel import Panel
+from rich.syntax import Syntax
+from rich.table import Table
+from rich import print as rprint
+
+# Create a Rich console object
+console = Console()
 
 def send_receive(sock, request_dict):
     """Send a request to the server and receive the response."""
@@ -60,6 +68,26 @@ def receive_response(sock):
     response_json = b''.join(chunks).decode('utf-8')
     return json.loads(response_json)
 
+def display_response(response, title="Response"):
+    """Display a response using Rich formatting."""
+    # Create a table for the response
+    table = Table(title=title)
+    table.add_column("Key", style="cyan")
+    table.add_column("Value", style="green")
+
+    for key, value in response.items():
+        if key == "output":
+            # Format output as a panel
+            table.add_row(key, Panel(value.strip(), border_style="blue"))
+        elif key == "error":
+            # Format error with red text
+            table.add_row(key, f"[bold red]{value}[/bold red]")
+        else:
+            # Format other values normally
+            table.add_row(key, str(value))
+
+    console.print(table)
+
 def main():
     # Server connection details
     host = "localhost"
@@ -69,33 +97,34 @@ def main():
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     try:
         sock.connect((host, port))
-        print(f"Connected to {host}:{port}")
+        console.print(f"Connected to [bold]{host}:{port}[/bold]", style="green")
 
         # Receive initial greeting using the length-prefixed protocol
         greeting = receive_response(sock)
-        print(f"Server greeting: {json.dumps(greeting, indent=2)}")
+        console.print(Panel(Syntax(json.dumps(greeting, indent=2), "json", theme="monokai"), 
+                           title="Server Greeting", border_style="green"))
 
         # Test 1: Execute code without a session ID (creates a new session)
-        print("\n--- Test 1: Execute code without a session ID ---")
+        console.rule("[bold]Test 1: Execute code without a session ID[/bold]")
         response = send_receive(sock, {
             "code": "x = 42\nprint(f'x = {x}')"
         })
-        print(f"Response: {json.dumps(response, indent=2)}")
+        display_response(response, "Test 1 Response")
 
         # Save the session ID for later use
         session_id = response.get("session_id")
-        print(f"Session ID: {session_id}")
+        console.print(f"Session ID: [bold cyan]{session_id}[/bold cyan]")
 
         # Test 2: Execute code in the same session (using the session ID)
-        print("\n--- Test 2: Execute code in the same session ---")
+        console.rule("[bold]Test 2: Execute code in the same session[/bold]")
         response = send_receive(sock, {
             "code": "y = x * 2\nprint(f'y = {y}')",
             "session_id": session_id
         })
-        print(f"Response: {json.dumps(response, indent=2)}")
+        display_response(response, "Test 2 Response")
 
         # Test 3: Define a function in the session
-        print("\n--- Test 3: Define a function in the session ---")
+        console.rule("[bold]Test 3: Define a function in the session[/bold]")
         response = send_receive(sock, {
             "code": """
 def greet(name):
@@ -104,46 +133,46 @@ def greet(name):
 """,
             "session_id": session_id
         })
-        print(f"Response: {json.dumps(response, indent=2)}")
+        display_response(response, "Test 3 Response")
 
         # Test 4: Call the function defined in the previous request
-        print("\n--- Test 4: Call the function defined in the previous request ---")
+        console.rule("[bold]Test 4: Call the function defined in the previous request[/bold]")
         response = send_receive(sock, {
             "code": "print(greet('Python'))",
             "session_id": session_id
         })
-        print(f"Response: {json.dumps(response, indent=2)}")
+        display_response(response, "Test 4 Response")
 
         # Test 5: Create a new session
-        print("\n--- Test 5: Create a new session ---")
+        console.rule("[bold]Test 5: Create a new session[/bold]")
         response = send_receive(sock, {
             "code": "print('This is a new session')"
         })
-        print(f"Response: {json.dumps(response, indent=2)}")
+        display_response(response, "Test 5 Response")
         new_session_id = response.get("session_id")
-        print(f"New Session ID: {new_session_id}")
+        console.print(f"New Session ID: [bold cyan]{new_session_id}[/bold cyan]")
 
         # Test 6: Verify the new session doesn't have access to variables from the first session
-        print("\n--- Test 6: Verify session isolation ---")
+        console.rule("[bold]Test 6: Verify session isolation[/bold]")
         response = send_receive(sock, {
             "code": "try:\n    print(f'x = {x}')\nexcept NameError as e:\n    print(f'Error: {e}')",
             "session_id": new_session_id
         })
-        print(f"Response: {json.dumps(response, indent=2)}")
+        display_response(response, "Test 6 Response")
 
         # Test 7: Try to access a non-existent session
-        print("\n--- Test 7: Try to access a non-existent session ---")
+        console.rule("[bold]Test 7: Try to access a non-existent session[/bold]")
         response = send_receive(sock, {
             "code": "print('This should fail')",
             "session_id": "non-existent-session-id"
         })
-        print(f"Response: {json.dumps(response, indent=2)}")
+        display_response(response, "Test 7 Response")
 
     except Exception as e:
-        print(f"Error: {e}")
+        console.print(f"[bold red]Error:[/bold red] {e}")
     finally:
         sock.close()
-        print("Connection closed")
+        console.print("Connection closed", style="yellow")
 
 if __name__ == "__main__":
     main()