more-compute 0.2.6__py3-none-any.whl → 0.3.0__py3-none-any.whl

morecompute/utils/shell_utils.py

@@ -0,0 +1,68 @@
+"""
+Shared utilities for shell command execution.
+Used by both special_commands.py and cell_magics.py to avoid duplication.
+"""
+import os
+import shlex
+import platform
+from typing import List, Dict
+
+
+def prepare_shell_command(cmd: str) -> List[str]:
+    """
+    Convert shell command to subprocess-compatible format.
+    Handles pip routing, python unbuffering, and platform-specific shells.
+
+    Args:
+        cmd: Shell command string (e.g., "pip install pandas")
+
+    Returns:
+        List of command arguments for subprocess
+    """
+    if cmd.startswith('pip '):
+        # Route pip through Python module for better control
+        parts = ['python', '-m'] + shlex.split(cmd)
+        # Add progress bar control for pip
+        if 'install' in cmd and '--progress-bar' not in cmd:
+            parts.extend(['--progress-bar', 'off'])
+        return parts
+
+    elif cmd.startswith('python '):
+        # Add unbuffered flag to python commands
+        parts = shlex.split(cmd)
+        parts.insert(1, '-u')  # Add -u after 'python'
+        return parts
+
+    else:
+        # For other shell commands, use platform-appropriate shell
+        system = platform.system()
+        if system == 'Windows':
+            return ['cmd', '/c', cmd]
+        elif system == 'Darwin':
+            return ['/bin/bash', '-c', cmd]
+        else:
+            return ['/bin/bash', '-c', cmd]
+
+
+def prepare_shell_environment(cmd: str) -> Dict[str, str]:
+    """
+    Prepare environment variables for shell command execution.
+
+    Args:
+        cmd: Shell command string
+
+    Returns:
+        Environment dictionary
+    """
+    env = os.environ.copy()
+
+    # Always set unbuffered Python
+    env['PYTHONUNBUFFERED'] = '1'
+    env['PYTHONDONTWRITEBYTECODE'] = '1'
+
+    # Additional settings for pip commands
+    if 'pip install' in cmd:
+        env['PIP_DISABLE_PIP_VERSION_CHECK'] = '1'
+        env['PIP_NO_CACHE_DIR'] = '1'
+
+    return env

morecompute/utils/special_commands.py

@@ -5,10 +5,15 @@ import asyncio
 import subprocess
 import time
 import shlex
+import platform
 from contextlib import redirect_stdout, redirect_stderr
 from typing import Dict, Any, Optional, Tuple, Union
 from fastapi import WebSocket
 
+from .cell_magics import CellMagicHandlers
+from .line_magics import LineMagicHandlers
+from .shell_utils import prepare_shell_command, prepare_shell_environment
+
 
 # this file is not tested that all functions work, need to write a test file / manually check
 # to-do
@@ -19,14 +24,28 @@ class AsyncSpecialCommandHandler:
     def __init__(self, globals_dict: dict):
         self.globals_dict = globals_dict
         self.captured_outputs = {}  # Store captured outputs from %%capture
+        self.cell_magic_handlers = CellMagicHandlers(globals_dict, self)
+        self.line_magic_handlers = LineMagicHandlers(globals_dict)
 
     def is_special_command(self, source_code: Union[str, list, tuple]) -> bool:
-        """Check if the source code is a special command"""
+        """Check if the source code is a special command or contains shell commands"""
         text = self._coerce_source_to_text(source_code)
         stripped = text.strip()
-        return (stripped.startswith('!') or
-                stripped.startswith('%%') or
-                stripped.startswith('%'))
+
+        # Check if starts with magic or shell command
+        if (stripped.startswith('!') or
+            stripped.startswith('%%') or
+            stripped.startswith('%')):
+            return True
+
+        # Check if ANY line contains a shell command (like Jupyter/Colab)
+        # This allows mixing Python code with !commands
+        lines = text.split('\n')
+        for line in lines:
+            if line.strip().startswith('!'):
+                return True
+
+        return False
 
     async def execute_special_command(self, source_code: Union[str, list, tuple], result: Dict[str, Any],
                                     start_time: float, execution_count: int,
@@ -42,8 +61,27 @@ class AsyncSpecialCommandHandler:
             return await self._execute_cell_magic(text, result, start_time, execution_count, websocket)
         elif stripped.startswith('%'):
             return await self._execute_line_magic(stripped[1:], result, start_time, websocket)
-        else:
-            raise ValueError("Not a special command")
+
+        # Cell contains shell commands mixed with Python code
+        # Treat it like regular code execution but preprocess shell commands
+        # This reuses the preprocessing logic from cell_magics.py
+        try:
+            stdout_text, stderr_text = await self.cell_magic_handlers.execute_cell_content(
+                text, result, execution_count, websocket,
+                capture_stdout=False, capture_stderr=False
+            )
+            result["status"] = "ok"
+        except Exception as e:
+            result["status"] = "error"
+            result["error"] = {
+                "ename": type(e).__name__,
+                "evalue": str(e),
+                "traceback": [f"Error executing cell: {str(e)}"]
+            }
+
+        # Calculate execution time
+        result["execution_time"] = f"{(time.time() - start_time) * 1000:.1f}ms"
+        return result
 
     def _coerce_source_to_text(self, source_code: Union[str, list, tuple]) -> str:
         """Normalize incoming source to a single text string"""
@@ -61,9 +99,9 @@ class AsyncSpecialCommandHandler:
                                    cell_index: Optional[int] = None) -> Dict[str, Any]:
         """Execute a shell command with real-time streaming output"""
         try:
-            # Prepare environment and command for streaming
-            env = self._prepare_streaming_environment(command)
-            cmd_parts = self._prepare_command_parts(command)
+            # Prepare environment and command for streaming (using shared utilities)
+            env = prepare_shell_environment(command)
+            cmd_parts = prepare_shell_command(command)
 
             # Send execution start notification
             if websocket:
@@ -112,6 +150,8 @@ class AsyncSpecialCommandHandler:
             # If pip install/uninstall occurred, notify clients to refresh packages
             try:
                 if websocket and (command.startswith('pip install') or command.startswith('pip uninstall') or 'pip install' in command or 'pip uninstall' in command):
+                    # Small delay to ensure pip finishes writing metadata to disk
+                    await asyncio.sleep(0.5)
                     await websocket.send_json({
                         "type": "packages_updated",
                         "data": {"action": "pip"}
@@ -158,57 +198,21 @@ class AsyncSpecialCommandHandler:
         # Placeholder for future process-based interruption logic
         return
 
-    def _prepare_streaming_environment(self, shell_cmd: str) -> dict:
-        """Prepare environment variables for unbuffered output"""
-        env = os.environ.copy()
-
-        # Always set unbuffered Python
-        env['PYTHONUNBUFFERED'] = '1'
-        env['PYTHONDONTWRITEBYTECODE'] = '1'
-
-        # Additional settings for specific commands
-        if 'pip install' in shell_cmd:
-            env['PIP_DISABLE_PIP_VERSION_CHECK'] = '1'
-            env['PIP_NO_CACHE_DIR'] = '1'
-
-        return env
-
-    def _prepare_command_parts(self, shell_cmd: str) -> list:
-        """Convert shell command to subprocess-compatible format"""
-
-        if shell_cmd.startswith('pip '):
-            # Route pip through Python module for better control
-            parts = ['python', '-m'] + shlex.split(shell_cmd)
-            # Add progress bar control for pip
-            if 'install' in shell_cmd and '--progress-bar' not in shell_cmd:
-                parts.extend(['--progress-bar', 'off'])
-            return parts
-
-        elif shell_cmd.startswith('python '):
-            # Add unbuffered flag to python commands
-            parts = shlex.split(shell_cmd)
-            parts.insert(1, '-u')  # Add -u after 'python'
-            return parts
-
-        else:
-            # For other shell commands, use shell execution
-            return ['/bin/zsh', '-c', shell_cmd]  # macOS with zsh
-
     async def _stream_output(self, stream, stream_type: str, result: Dict[str, Any],
                            websocket: Optional[WebSocket] = None,
                            cell_index: Optional[int] = None):
         """Read from a stream and send to websocket, while capturing the output."""
-        
+
         output_text = ""
         while True:
             try:
                 line = await stream.readline()
                 if not line:
                     break
-                
+
                 decoded_line = line.decode('utf-8')
                 output_text += decoded_line
-                
+
                 if websocket:
                     await websocket.send_json({
                         "type": "stream_output",
@@ -234,7 +238,7 @@ class AsyncSpecialCommandHandler:
                         }
                     })
                 break
-        
+
         # Add the captured text to the final result object
         if output_text:
             # Look for an existing stream output of the same type to append to
@@ -261,13 +265,37 @@ class AsyncSpecialCommandHandler:
         magic_name = magic_parts[0][2:]  # Remove %%
         magic_args = magic_parts[1:] if len(magic_parts) > 1 else []
 
+        # Map magic names to handler methods
+        magic_handlers = {
+            "capture": lambda: self.cell_magic_handlers.handle_capture(
+                magic_args, cell_content, result, start_time, execution_count, websocket
+            ),
+            "time": lambda: self.cell_magic_handlers.handle_time(
+                cell_content, result, start_time, execution_count, websocket
+            ),
+            "timeit": lambda: self.cell_magic_handlers.handle_timeit(
+                magic_args, cell_content, result, start_time, execution_count, websocket
+            ),
+            "writefile": lambda: self.cell_magic_handlers.handle_writefile(
+                magic_args, cell_content, result, start_time, websocket
+            ),
+            "bash": lambda: self.cell_magic_handlers.handle_bash(
+                cell_content, result, start_time, websocket
+            ),
+            "sh": lambda: self.cell_magic_handlers.handle_bash(
+                cell_content, result, start_time, websocket
+            ),
+            "html": lambda: self.cell_magic_handlers.handle_html(
+                cell_content, result, start_time, websocket
+            ),
+            "markdown": lambda: self.cell_magic_handlers.handle_markdown(
+                cell_content, result, start_time, websocket
+            ),
+        }
+
         try:
-            if magic_name == "capture":
-                return await self._handle_capture_magic(magic_args, cell_content, result, start_time, execution_count, websocket)
-            elif magic_name == "time":
-                return await self._handle_time_magic(cell_content, result, start_time, execution_count, websocket)
-            elif magic_name == "writefile":
-                return await self._handle_writefile_magic(magic_args, cell_content, result, start_time, websocket)
+            if magic_name in magic_handlers:
+                return await magic_handlers[magic_name]()
             else:
                 result["status"] = "error"
                 result["error"] = {
@@ -310,19 +338,30 @@ class AsyncSpecialCommandHandler:
         magic_name = parts[0]
         magic_args = parts[1:] if len(parts) > 1 else []
 
+        # Map magic names to handler methods
+        magic_handlers = {
+            "pwd": lambda: self.line_magic_handlers.handle_pwd(magic_args, result, websocket),
+            "cd": lambda: self.line_magic_handlers.handle_cd(magic_args, result, websocket),
+            "ls": lambda: self.line_magic_handlers.handle_ls(magic_args, result, websocket),
+            "env": lambda: self.line_magic_handlers.handle_env(magic_args, result, websocket),
+            "who": lambda: self.line_magic_handlers.handle_who(magic_args, result, websocket),
+            "whos": lambda: self.line_magic_handlers.handle_whos(magic_args, result, websocket),
+            "time": lambda: self.line_magic_handlers.handle_time(magic_args, result, websocket),
+            "timeit": lambda: self.line_magic_handlers.handle_timeit(magic_args, result, websocket),
+            "pip": lambda: self.line_magic_handlers.handle_pip(magic_args, result, self, websocket),
+            "load": lambda: self.line_magic_handlers.handle_load(magic_args, result, websocket),
+            "reset": lambda: self.line_magic_handlers.handle_reset(magic_args, result, websocket),
+            "lsmagic": lambda: self.line_magic_handlers.handle_lsmagic(magic_args, result, websocket),
+            "matplotlib": lambda: self.line_magic_handlers.handle_matplotlib(magic_args, result, websocket),
+            "load_ext": lambda: self.line_magic_handlers.handle_load_ext(magic_args, result, websocket),
+            "reload_ext": lambda: self.line_magic_handlers.handle_reload_ext(magic_args, result, websocket),
+            "unload_ext": lambda: self.line_magic_handlers.handle_unload_ext(magic_args, result, websocket),
+            "run": lambda: self.line_magic_handlers.handle_run(magic_args, result, websocket),
+        }
+
         try:
-            if magic_name == "pwd":
-                return await self._handle_pwd_magic(result, start_time, websocket)
-            elif magic_name == "cd":
-                return await self._handle_cd_magic(magic_args, result, start_time, websocket)
-            elif magic_name == "ls":
-                return await self._handle_ls_magic(magic_args, result, start_time, websocket)
-            elif magic_name == "env":
-                return await self._handle_env_magic(magic_args, result, start_time, websocket)
-            elif magic_name == "who":
-                return await self._handle_who_magic(result, start_time, websocket)
-            elif magic_name == "whos":
-                return await self._handle_whos_magic(result, start_time, websocket)
+            if magic_name in magic_handlers:
+                return await magic_handlers[magic_name]()
             else:
                 result["status"] = "error"
                 result["error"] = {
@@ -356,109 +395,3 @@ class AsyncSpecialCommandHandler:
 
         result["execution_time"] = f"{(time.time() - start_time) * 1000:.1f}ms"
         return result
-
-    # Cell Magic Implementations
-
-    async def _handle_capture_magic(self, args: list, cell_content: str, result: Dict[str, Any],
-                                  start_time: float, execution_count: int,
-                                  websocket: Optional[WebSocket] = None) -> Dict[str, Any]:
-        """Handle %%capture magic - capture stdout/stderr without displaying"""
-        output_var = args[0] if args else None
-        no_stdout = "--no-stdout" in args
-        no_stderr = "--no-stderr" in args
-
-        # Capture outputs
-        stdout_capture = None if no_stdout else io.StringIO()
-        stderr_capture = None if no_stderr else io.StringIO()
-
-        try:
-            # Execute the cell content with output capture
-            if stdout_capture or stderr_capture:
-                with redirect_stdout(stdout_capture or sys.stdout), \
-                     redirect_stderr(stderr_capture or sys.stderr):
-                    compiled_code = compile(cell_content, '<cell>', 'exec')
-                    exec(compiled_code, self.globals_dict)
-            else:
-                compiled_code = compile(cell_content, '<cell>', 'exec')
-                exec(compiled_code, self.globals_dict)
-
-            # Store captured output in a variable if specified
-            if output_var:
-                captured_data = {
-                    'stdout': stdout_capture.getvalue() if stdout_capture else '',
-                    'stderr': stderr_capture.getvalue() if stderr_capture else ''
-                }
-                self.globals_dict[output_var] = captured_data
-                self.captured_outputs[output_var] = captured_data
-
-            # Don't add outputs to result (they're captured, not displayed)
-            if websocket:
-                await websocket.send_json({
-                    "type": "execution_complete",
-                    "data": {
-                        "status": "ok",
-                        "message": "Output captured" + (f" in variable '{output_var}'" if output_var else "")
-                    }
-                })
-
-        except Exception as e:
-            result["status"] = "error"
-            result["error"] = {
-                "ename": type(e).__name__,
-                "evalue": str(e),
-                "traceback": [f"Capture magic error: {str(e)}"]
-            }
-
-            if websocket:
-                await websocket.send_json({
-                    "type": "execution_error",
-                    "data": {
-                        "error": result["error"]
-                    }
-                })
-
-        return result
-
-    # Add other magic method implementations here...
-    # (Time magic, writefile magic, line magics like pwd, cd, ls, etc.)
-    # I'll implement a few key ones to keep this focused:
-
-    async def _handle_pwd_magic(self, result: Dict[str, Any], start_time: float,
-                              websocket: Optional[WebSocket] = None) -> Dict[str, Any]:
-        """Handle %pwd magic - print working directory"""
-        try:
-            pwd = os.getcwd()
-            output_data = {
-                "output_type": "execute_result",
-                "execution_count": None,
-                "data": {
-                    "text/plain": f"'{pwd}'"
-                }
-            }
-            result["outputs"].append(output_data)
-
-            if websocket:
-                await websocket.send_json({
-                    "type": "execute_result",
-                    "data": {
-                        "data": output_data["data"]
-                    }
-                })
-
-        except Exception as e:
-            result["status"] = "error"
-            result["error"] = {
-                "ename": type(e).__name__,
-                "evalue": str(e),
-                "traceback": [f"PWD magic error: {str(e)}"]
-            }
-
-            if websocket:
-                await websocket.send_json({
-                    "type": "execution_error",
-                    "data": {
-                        "error": result["error"]
-                    }
-                })
-
-        return result