more-compute 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

morecompute/notebook.py

@@ -1,6 +1,8 @@
 import json
+from pathlib import Path
 from typing import List, Dict, Any
 from uuid import uuid4
+from .utils.py_percent_parser import parse_py_percent, generate_py_percent
 
 class Notebook:
     """Manages the state of a notebook's cells."""
@@ -71,6 +73,8 @@ class Notebook:
             if cell.get('cell_type') == 'code':
                 cell['outputs'] = []
                 cell['execution_count'] = None
+                cell['execution_time'] = None
+                cell['error'] = None
 
     def to_json(self) -> str:
         # Basic notebook format
@@ -84,10 +88,18 @@ class Notebook:
 
     def load_from_file(self, file_path: str):
         try:
-            with open(file_path, 'r') as f:
-                data = json.load(f)
+            path = Path(file_path)
+
+            # Check file extension
+            if path.suffix == '.py':
+                # Load .py file with py:percent format
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+
+                data = parse_py_percent(content)
                 loaded_cells = data.get('cells', [])
-                # Ensure stable IDs for all cells (back-compat for notebooks without IDs)
+
+                # Ensure stable IDs for all cells
                 self.cells = []
                 for cell in loaded_cells:
                     if not isinstance(cell, dict):
@@ -95,9 +107,40 @@ class Notebook:
                     if 'id' not in cell or not cell['id']:
                         cell['id'] = self._generate_cell_id()
                     self.cells.append(cell)
+
                 self.metadata = data.get('metadata', {})
                 self.file_path = file_path
-        except (FileNotFoundError, json.JSONDecodeError) as e:
+
+            elif path.suffix == '.ipynb':
+                # Block .ipynb files with helpful error
+                raise ValueError(
+                    f"MoreCompute only supports .py notebooks.\n\n"
+                    f"Convert your notebook with:\n"
+                    f"  more-compute convert {path.name} -o {path.stem}.py\n\n"
+                    f"Then open with:\n"
+                    f"  more-compute {path.stem}.py"
+                )
+
+            else:
+                raise ValueError(f"Unsupported file format: {path.suffix}. Use .py files.")
+
+        except FileNotFoundError as e:
+            print(f"Error: File not found: {e}")
+            # Initialize with a default cell if loading fails
+            self.cells = [{
+                'id': self._generate_cell_id(),
+                'cell_type': 'code',
+                'source': '',
+                'metadata': {},
+                'outputs': [],
+                'execution_count': None
+            }]
+            self.metadata = {}
+            self.file_path = file_path
+        except ValueError as e:
+            # Re-raise validation errors (like .ipynb block)
+            raise
+        except Exception as e:
             print(f"Error loading notebook: {e}")
             # Initialize with a default cell if loading fails
             self.cells = [{
@@ -116,8 +159,24 @@ class Notebook:
         if not path_to_save:
             raise ValueError("No file path specified for saving.")
 
-        with open(path_to_save, 'w') as f:
-            f.write(self.to_json())
+        path = Path(path_to_save)
+
+        # Save in appropriate format based on extension
+        if path.suffix == '.py':
+            # Save as py:percent format
+            content = generate_py_percent(self.cells)
+            with open(path_to_save, 'w', encoding='utf-8') as f:
+                f.write(content)
+        elif path.suffix == '.ipynb':
+            # Block saving as .ipynb
+            raise ValueError("MoreCompute only supports .py notebooks. Use .py extension.")
+        else:
+            # Default to .py if no extension
+            path_to_save = str(path.with_suffix('.py'))
+            content = generate_py_percent(self.cells)
+            with open(path_to_save, 'w', encoding='utf-8') as f:
+                f.write(content)
+
         self.file_path = path_to_save
 
     def _generate_cell_id(self) -> str:

morecompute/server.py

@@ -99,7 +99,23 @@ async def startup_event():
 @app.on_event("shutdown")
 async def shutdown_event():
     """Cleanup services on shutdown."""
-    global lsp_service
+    global lsp_service, executor
+
+    # Shutdown executor and worker process
+    if executor and executor.worker_proc:
+        try:
+            print("[EXECUTOR] Shutting down worker process...", file=sys.stderr, flush=True)
+            executor.worker_proc.terminate()
+            executor.worker_proc.wait(timeout=2)
+            print("[EXECUTOR] Worker process shutdown complete", file=sys.stderr, flush=True)
+        except Exception as e:
+            print(f"[EXECUTOR] Error during worker shutdown, forcing kill: {e}", file=sys.stderr, flush=True)
+            try:
+                executor.worker_proc.kill()
+            except Exception:
+                pass
+
+    # Shutdown LSP service
     if lsp_service:
         try:
             await lsp_service.shutdown()
@@ -441,12 +457,34 @@ class WebSocketManager:
 
     async def handle_message_loop(self, websocket: WebSocket):
         """Main loop to handle incoming WebSocket messages."""
+        tasks = set()
+
+        def task_done_callback(task):
+            tasks.discard(task)
+            # Check for exceptions in completed tasks
+            try:
+                exc = task.exception()
+                if exc:
+                    print(f"[SERVER] Task raised exception: {exc}", file=sys.stderr, flush=True)
+                    import traceback
+                    traceback.print_exception(type(exc), exc, exc.__traceback__)
+            except asyncio.CancelledError:
+                pass
+            except Exception as e:
+                print(f"[SERVER] Error in task_done_callback: {e}", file=sys.stderr, flush=True)
+
         while True:
             try:
                 message = await websocket.receive_json()
-                await self._handle_message(websocket, message)
+                # Process messages concurrently so interrupts can arrive during execution
+                task = asyncio.create_task(self._handle_message(websocket, message))
+                tasks.add(task)
+                task.add_done_callback(task_done_callback)
             except WebSocketDisconnect:
                 self.disconnect(websocket)
+                # Cancel all pending tasks
+                for task in tasks:
+                    task.cancel()
                 break
             except Exception as e:
                 await self._send_error(websocket, f"Unhandled error: {e}")
@@ -534,12 +572,24 @@ class WebSocketManager:
         else:
             # Normal add cell
             self.notebook.add_cell(index=index, cell_type=cell_type, source=source)
+
+        # Save the notebook after adding cell
+        try:
+            self.notebook.save_to_file()
+        except Exception as e:
+            print(f"Warning: Failed to save notebook after adding cell: {e}", file=sys.stderr)
+
         await self.broadcast_notebook_update()
 
     async def _handle_delete_cell(self, websocket: WebSocket, data: dict):
         index = data.get('cell_index')
         if index is not None:
             self.notebook.delete_cell(index)
+            # Save the notebook after deleting cell
+            try:
+                self.notebook.save_to_file()
+            except Exception as e:
+                print(f"Warning: Failed to save notebook after deleting cell: {e}", file=sys.stderr)
             await self.broadcast_notebook_update()
 
     async def _handle_update_cell(self, websocket: WebSocket, data: dict):
@@ -587,49 +637,30 @@ class WebSocketManager:
         print(f"[SERVER] Interrupt request received for cell {cell_index}", file=sys.stderr, flush=True)
 
         # Perform the interrupt (this may take up to 1 second)
+        # The execution handler will send the appropriate error and completion messages
         await self.executor.interrupt_kernel(cell_index=cell_index)
 
-        print(f"[SERVER] Interrupt completed, sending error message", file=sys.stderr, flush=True)
+        print(f"[SERVER] Interrupt completed, execution handler will send completion messages", file=sys.stderr, flush=True)
 
-        # Inform all clients that the currently running cell (if any) is interrupted
-        try:
-            await websocket.send_json({
-                "type": "execution_error",
-                "data": {
-                    "cell_index": cell_index,
-                    "error": {
-                        "output_type": "error",
-                        "ename": "KeyboardInterrupt",
-                        "evalue": "Execution interrupted by user",
-                        "traceback": ["KeyboardInterrupt: Execution was stopped by user"]
-                    }
-                }
-            })
-            await websocket.send_json({
-                "type": "execution_complete",
-                "data": {
-                    "cell_index": cell_index,
-                    "result": {
-                        "status": "error",
-                        "execution_count": None,
-                        "execution_time": "interrupted",
-                        "outputs": [],
-                        "error": {
-                            "output_type": "error",
-                            "ename": "KeyboardInterrupt",
-                            "evalue": "Execution interrupted by user",
-                            "traceback": ["KeyboardInterrupt: Execution was stopped by user"]
-                        }
-                    }
-                }
-            })
-            print(f"[SERVER] Error messages sent for cell {cell_index}", file=sys.stderr, flush=True)
-        except Exception as e:
-            print(f"[SERVER] Failed to send error messages: {e}", file=sys.stderr, flush=True)
+        # Note: We don't send completion messages here anymore because:
+        # 1. For shell commands: AsyncSpecialCommandHandler._execute_shell_command sends them
+        # 2. For Python code: The worker sends them
+        # Sending duplicate messages causes the frontend to get confused
 
     async def _handle_reset_kernel(self, websocket: WebSocket, data: dict):
+        import sys
+        print(f"[SERVER] Resetting kernel", file=sys.stderr, flush=True)
         self.executor.reset_kernel()
         self.notebook.clear_all_outputs()
+
+        # Note: We don't save the notebook here - this preserves execution times
+        # from the last session, which is useful for seeing how long things took
+
+        # Broadcast kernel restart to all clients
+        await self.broadcast_pod_update({
+            "type": "kernel_restarted",
+            "data": {}
+        })
         await self.broadcast_notebook_update()
 
     async def _send_error(self, websocket: WebSocket, error_message: str):
@@ -811,9 +842,10 @@ async def _connect_to_pod_background(pod_id: str):
             reconnect_zmq_sockets(
                 executor,
                 cmd_addr=addresses["cmd_addr"],
-                pub_addr=addresses["pub_addr"]
+                pub_addr=addresses["pub_addr"],
+                is_remote=True  # Critical: Tell executor this is a remote worker
             )
-            print(f"[CONNECT BACKGROUND] Successfully connected to pod {pod_id}", file=sys.stderr, flush=True)
+            print(f"[CONNECT BACKGROUND] Successfully connected to pod {pod_id}, executor.is_remote=True", file=sys.stderr, flush=True)
         else:
             # Connection failed - clean up
             print(f"[CONNECT BACKGROUND] Failed to connect: {result}", file=sys.stderr, flush=True)

morecompute/utils/cell_magics.py

@@ -248,6 +248,11 @@ class CellMagicHandlers:
                     env=env
                 )
 
+                # Track process for interrupt handling
+                if hasattr(cell_magic_handler, 'special_handler'):
+                    cell_magic_handler.special_handler.current_process_sync = process
+                    print(f"[CELL_MAGIC] Tracking sync subprocess PID={process.pid}", file=sys.stderr, flush=True)
+
                 # Read and print output line by line (real-time streaming)
                 def read_stream(stream, output_type):
                     """Read stream line by line and print immediately"""
@@ -276,12 +281,38 @@ class CellMagicHandlers:
                 stdout_thread.start()
                 stderr_thread.start()
 
-                # Wait for process to complete
-                return_code = process.wait()
+                # Wait for process to complete, checking if it was killed
+                try:
+                    # Poll with timeout to detect if process was killed externally
+                    while process.poll() is None:
+                        try:
+                            process.wait(timeout=0.1)
+                        except subprocess.TimeoutExpired:
+                            # Check if interrupted
+                            if hasattr(cell_magic_handler, 'special_handler'):
+                                if cell_magic_handler.special_handler.sync_interrupted:
+                                    # Process was killed by interrupt handler
+                                    print(f"[CELL_MAGIC] Process was interrupted, raising KeyboardInterrupt", file=sys.stderr, flush=True)
+                                    raise KeyboardInterrupt("Execution interrupted by user")
+
+                    return_code = process.returncode
+                except KeyboardInterrupt:
+                    # Kill process if KeyboardInterrupt
+                    try:
+                        process.kill()
+                        process.wait()
+                    except Exception:
+                        pass
+                    raise
 
                 # Wait for output threads to finish
-                stdout_thread.join()
-                stderr_thread.join()
+                stdout_thread.join(timeout=1)
+                stderr_thread.join(timeout=1)
+
+                # Clear process reference
+                if hasattr(cell_magic_handler, 'special_handler'):
+                    cell_magic_handler.special_handler.current_process_sync = None
+                    print(f"[CELL_MAGIC] Cleared sync subprocess reference", file=sys.stderr, flush=True)
 
                 return return_code
 

morecompute/utils/notebook_converter.py

@@ -0,0 +1,129 @@
+"""Converter utilities for notebook formats."""
+
+import json
+import re
+from pathlib import Path
+from typing import List, Set
+from .py_percent_parser import generate_py_percent, parse_py_percent
+
+
+def extract_pip_dependencies(notebook_data: dict) -> Set[str]:
+    """
+    Extract package names from !pip install and %pip install commands.
+
+    Args:
+        notebook_data: Parsed notebook JSON
+
+    Returns:
+        Set of package names
+    """
+    packages = set()
+
+    for cell in notebook_data.get('cells', []):
+        if cell.get('cell_type') != 'code':
+            continue
+
+        source = cell.get('source', [])
+        if isinstance(source, list):
+            source = ''.join(source)
+
+        # Match: !pip install package1 package2
+        # Match: %pip install package1 package2
+        pip_pattern = r'[!%]pip\s+install\s+([^\n]+)'
+        matches = re.finditer(pip_pattern, source)
+
+        for match in matches:
+            install_line = match.group(1)
+            # Remove common flags
+            install_line = re.sub(r'--[^\s]+\s*', '', install_line)
+            install_line = re.sub(r'-[qU]\s*', '', install_line)
+
+            # Extract package names (handle package==version format)
+            parts = install_line.split()
+            for part in parts:
+                part = part.strip()
+                if part and not part.startswith('-'):
+                    packages.add(part)
+
+    return packages
+
+
+def convert_ipynb_to_py(ipynb_path: Path, output_path: Path, include_uv_deps: bool = True) -> None:
+    """
+    Convert .ipynb notebook to .py format with py:percent cell markers.
+
+    Args:
+        ipynb_path: Path to input .ipynb file
+        output_path: Path to output .py file
+        include_uv_deps: Whether to add UV inline script dependencies
+    """
+    # Read notebook
+    with open(ipynb_path, 'r', encoding='utf-8') as f:
+        notebook_data = json.load(f)
+
+    cells = notebook_data.get('cells', [])
+
+    # Generate UV dependencies header if requested
+    header_lines = []
+    if include_uv_deps:
+        dependencies = extract_pip_dependencies(notebook_data)
+        if dependencies:
+            header_lines.append('# /// script')
+            header_lines.append('# dependencies = [')
+            for dep in sorted(dependencies):
+                header_lines.append(f'#   "{dep}",')
+            header_lines.append('# ]')
+            header_lines.append('# ///')
+            header_lines.append('')
+
+    # Generate py:percent format
+    py_content = generate_py_percent(cells)
+
+    # Combine header and content
+    if header_lines:
+        final_content = '\n'.join(header_lines) + '\n' + py_content
+    else:
+        final_content = py_content
+
+    # Write output
+    with open(output_path, 'w', encoding='utf-8') as f:
+        f.write(final_content)
+
+    print(f"✓ Converted {ipynb_path.name} → {output_path.name}")
+
+    # Show dependencies if found
+    if include_uv_deps and dependencies:
+        print(f"  Found dependencies: {', '.join(sorted(dependencies))}")
+        print(f"  Run with: more-compute {output_path.name}")
+
+
+def convert_py_to_ipynb(py_path: Path, output_path: Path) -> None:
+    """
+    Convert .py notebook to .ipynb format.
+
+    Args:
+        py_path: Path to input .py file
+        output_path: Path to output .ipynb file
+    """
+    # Read .py file
+    with open(py_path, 'r', encoding='utf-8') as f:
+        py_content = f.read()
+
+    # Parse py:percent format to notebook structure
+    notebook_data = parse_py_percent(py_content)
+
+    # Ensure source is in list format (Jupyter notebook standard)
+    for cell in notebook_data.get('cells', []):
+        source = cell.get('source', '')
+        if isinstance(source, str):
+            # Split into lines and keep newlines (Jupyter format)
+            lines = source.split('\n')
+            # Add \n to each line except the last
+            cell['source'] = [line + '\n' for line in lines[:-1]] + ([lines[-1]] if lines[-1] else [])
+
+    # Write .ipynb file
+    with open(output_path, 'w', encoding='utf-8') as f:
+        json.dump(notebook_data, f, indent=1, ensure_ascii=False)
+
+    print(f"Converted {py_path.name} -> {output_path.name}")
+    print(f"  Upload to Google Colab or open in Jupyter")

morecompute/utils/py_percent_parser.py

@@ -0,0 +1,190 @@
+"""Parser for py:percent format (.py files with # %% cell markers)."""
+
+import re
+from typing import List, Dict, Optional
+
+
+def parse_py_percent(content: str) -> Dict:
+    """
+    Parse py:percent format Python file into notebook structure.
+
+    Format:
+        # %%
+        code cell content
+
+        # %% [markdown]
+        # markdown content
+
+    Args:
+        content: Raw .py file content
+
+    Returns:
+        Dict with 'cells' and 'metadata' (compatible with Notebook class)
+    """
+    cells = []
+
+    # Split by cell markers (# %%)
+    # Keep the marker in the split to determine cell type
+    parts = re.split(r'(# %%.*?\n)', content)
+
+    # First part before any cell marker (usually imports/metadata)
+    if parts[0].strip():
+        # Check if it's UV inline script metadata
+        first_part = parts[0].strip()
+        if not first_part.startswith('# /// script'):
+            # It's code, add as first cell
+            cells.append({
+                'cell_type': 'code',
+                'source': parts[0],
+                'metadata': {},
+                'execution_count': None,
+                'outputs': []
+            })
+
+    # Process remaining parts
+    i = 1
+    while i < len(parts):
+        if i >= len(parts):
+            break
+
+        marker = parts[i] if i < len(parts) else ''
+        cell_content = parts[i + 1] if i + 1 < len(parts) else ''
+
+        # Determine cell type from marker
+        if '[markdown]' in marker.lower():
+            cell_type = 'markdown'
+            # Remove leading # from markdown lines
+            lines = cell_content.split('\n')
+            cleaned_lines = []
+            for line in lines:
+                if line.strip().startswith('#'):
+                    # Remove first # and any space after it
+                    cleaned_lines.append(line.strip()[1:].lstrip())
+                else:
+                    cleaned_lines.append(line)
+            cell_content = '\n'.join(cleaned_lines)
+        else:
+            cell_type = 'code'
+
+        # Only add non-empty cells
+        if cell_content.strip():
+            cell = {
+                'cell_type': cell_type,
+                'source': cell_content.strip(),
+                'metadata': {}
+            }
+
+            if cell_type == 'code':
+                cell['execution_count'] = None
+                cell['outputs'] = []
+
+            cells.append(cell)
+
+        i += 2
+
+    return {
+        'cells': cells,
+        'metadata': {
+            'kernelspec': {
+                'display_name': 'Python 3',
+                'language': 'python',
+                'name': 'python3'
+            },
+            'language_info': {
+                'name': 'python',
+                'version': '3.8.0'
+            }
+        },
+        'nbformat': 4,
+        'nbformat_minor': 4
+    }
+
+
+def extract_uv_dependencies(content: str) -> Optional[List[str]]:
+    """
+    Extract UV inline script dependencies from .py file.
+
+    Format:
+        # /// script
+        # dependencies = [
+        #   "package1",
+        #   "package2>=1.0.0",
+        # ]
+        # ///
+
+    Args:
+        content: Raw .py file content
+
+    Returns:
+        List of dependency strings, or None if no dependencies found
+    """
+    # Match UV inline script metadata block
+    pattern = r'# /// script\n(.*?)# ///\n'
+    match = re.search(pattern, content, re.DOTALL)
+
+    if not match:
+        return None
+
+    metadata_block = match.group(1)
+
+    # Extract dependencies list
+    deps_pattern = r'# dependencies = \[(.*?)\]'
+    deps_match = re.search(deps_pattern, metadata_block, re.DOTALL)
+
+    if not deps_match:
+        return None
+
+    deps_str = deps_match.group(1)
+
+    # Parse individual dependencies
+    dependencies = []
+    for line in deps_str.split('\n'):
+        line = line.strip()
+        if line.startswith('#'):
+            line = line[1:].strip()
+        # Remove quotes and trailing comma
+        line = line.strip('"\'').strip(',').strip()
+        if line:
+            dependencies.append(line)
+
+    return dependencies if dependencies else None
+
+
+def generate_py_percent(cells: List[Dict]) -> str:
+    """
+    Generate py:percent format from cell list.
+
+    Args:
+        cells: List of cell dicts with 'cell_type' and 'source'
+
+    Returns:
+        String in py:percent format
+    """
+    lines = []
+
+    for i, cell in enumerate(cells):
+        cell_type = cell.get('cell_type', 'code')
+        source = cell.get('source', '')
+
+        # Ensure source is string
+        if isinstance(source, list):
+            source = ''.join(source)
+
+        # Add cell marker
+        if cell_type == 'markdown':
+            lines.append('# %% [markdown]')
+            # Add # prefix to each line of markdown
+            for line in source.split('\n'):
+                if line.strip():
+                    lines.append(f'# {line}')
+                else:
+                    lines.append('#')
+        else:
+            lines.append('# %%')
+            lines.append(source)
+
+        # Add blank line between cells
+        if i < len(cells) - 1:
+            lines.append('')
+
+    return '\n'.join(lines)