feat: initial commit from workspace-mcp

gdocs/managers/table_operation_manager.py

@@ -0,0 +1,405 @@
+"""
+Table Operation Manager
+
+This module provides high-level table operations that orchestrate
+multiple Google Docs API calls for complex table manipulations.
+"""
+
+import logging
+import asyncio
+from typing import List, Dict, Any, Tuple
+
+from gdocs.docs_helpers import create_insert_table_request
+from gdocs.docs_structure import find_tables
+from gdocs.docs_tables import validate_table_data
+
+logger = logging.getLogger(__name__)
+
+
+class TableOperationManager:
+    """
+    High-level manager for Google Docs table operations.
+
+    Handles complex multi-step table operations including:
+    - Creating tables with data population
+    - Populating existing tables
+    - Managing cell-by-cell operations with proper index refreshing
+    """
+
+    def __init__(self, service):
+        """
+        Initialize the table operation manager.
+
+        Args:
+            service: Google Docs API service instance
+        """
+        self.service = service
+
+    async def create_and_populate_table(
+        self,
+        document_id: str,
+        table_data: List[List[str]],
+        index: int,
+        bold_headers: bool = True,
+        tab_id: str = None,
+    ) -> Tuple[bool, str, Dict[str, Any]]:
+        """
+        Creates a table and populates it with data in a reliable multi-step process.
+
+        This method extracts the complex logic from create_table_with_data tool function.
+
+        Args:
+            document_id: ID of the document to update
+            table_data: 2D list of strings for table content
+            index: Position to insert the table
+            bold_headers: Whether to make the first row bold
+            tab_id: Optional tab ID for targeting a specific tab
+
+        Returns:
+            Tuple of (success, message, metadata)
+        """
+        logger.debug(
+            f"Creating table at index {index}, dimensions: {len(table_data)}x{len(table_data[0]) if table_data and len(table_data) > 0 else 0}"
+        )
+
+        # Validate input data
+        is_valid, error_msg = validate_table_data(table_data)
+        if not is_valid:
+            return False, f"Invalid table data: {error_msg}", {}
+
+        rows = len(table_data)
+        cols = len(table_data[0])
+
+        try:
+            # Step 1: Create empty table
+            await self._create_empty_table(document_id, index, rows, cols, tab_id)
+
+            # Step 2: Get fresh document structure to find actual cell positions
+            fresh_tables = await self._get_document_tables(document_id, tab_id)
+            if not fresh_tables:
+                return False, "Could not find table after creation", {}
+
+            # Step 3: Populate each cell with proper index refreshing
+            population_count = await self._populate_table_cells(
+                document_id, table_data, bold_headers, tab_id
+            )
+
+            metadata = {
+                "rows": rows,
+                "columns": cols,
+                "populated_cells": population_count,
+                "table_index": len(fresh_tables) - 1,
+            }
+
+            return (
+                True,
+                f"Successfully created {rows}x{cols} table and populated {population_count} cells",
+                metadata,
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to create and populate table: {str(e)}")
+            return False, f"Table creation failed: {str(e)}", {}
+
+    async def _create_empty_table(
+        self, document_id: str, index: int, rows: int, cols: int, tab_id: str = None
+    ) -> None:
+        """Create an empty table at the specified index."""
+        logger.debug(f"Creating {rows}x{cols} table at index {index}")
+
+        await asyncio.to_thread(
+            self.service.documents()
+            .batchUpdate(
+                documentId=document_id,
+                body={
+                    "requests": [create_insert_table_request(index, rows, cols, tab_id)]
+                },
+            )
+            .execute
+        )
+
+    async def _get_document_tables(
+        self, document_id: str, tab_id: str = None
+    ) -> List[Dict[str, Any]]:
+        """Get fresh document structure and extract table information."""
+        doc = await asyncio.to_thread(
+            self.service.documents()
+            .get(documentId=document_id, includeTabsContent=True)
+            .execute
+        )
+
+        if tab_id:
+            tab = self._find_tab(doc.get("tabs", []), tab_id)
+            if tab and "documentTab" in tab:
+                doc = doc.copy()
+                doc["body"] = tab["documentTab"].get("body", {})
+
+        return find_tables(doc)
+
+    @staticmethod
+    def _find_tab(tabs: list, target_id: str):
+        """Recursively find a tab by ID."""
+        for tab in tabs:
+            if tab.get("tabProperties", {}).get("tabId") == target_id:
+                return tab
+            if "childTabs" in tab:
+                found = TableOperationManager._find_tab(tab["childTabs"], target_id)
+                if found:
+                    return found
+        return None
+
+    async def _populate_table_cells(
+        self,
+        document_id: str,
+        table_data: List[List[str]],
+        bold_headers: bool,
+        tab_id: str = None,
+    ) -> int:
+        """
+        Populate table cells with data, refreshing structure after each insertion.
+
+        This prevents index shifting issues by getting fresh cell positions
+        before each insertion.
+        """
+        population_count = 0
+
+        for row_idx, row_data in enumerate(table_data):
+            logger.debug(f"Processing row {row_idx}: {len(row_data)} cells")
+
+            for col_idx, cell_text in enumerate(row_data):
+                if not cell_text:  # Skip empty cells
+                    continue
+
+                try:
+                    # CRITICAL: Refresh document structure before each insertion
+                    success = await self._populate_single_cell(
+                        document_id,
+                        row_idx,
+                        col_idx,
+                        cell_text,
+                        bold_headers and row_idx == 0,
+                        tab_id,
+                    )
+
+                    if success:
+                        population_count += 1
+                        logger.debug(f"Populated cell ({row_idx},{col_idx})")
+                    else:
+                        logger.warning(f"Failed to populate cell ({row_idx},{col_idx})")
+
+                except Exception as e:
+                    logger.error(
+                        f"Error populating cell ({row_idx},{col_idx}): {str(e)}"
+                    )
+
+        return population_count
+
+    async def _populate_single_cell(
+        self,
+        document_id: str,
+        row_idx: int,
+        col_idx: int,
+        cell_text: str,
+        apply_bold: bool = False,
+        tab_id: str = None,
+    ) -> bool:
+        """
+        Populate a single cell with text, with optional bold formatting.
+
+        Returns True if successful, False otherwise.
+        """
+        try:
+            # Get fresh table structure to avoid index shifting issues
+            tables = await self._get_document_tables(document_id, tab_id)
+            if not tables:
+                return False
+
+            table = tables[-1]  # Use the last table (newly created one)
+            cells = table.get("cells", [])
+
+            # Bounds checking
+            if row_idx >= len(cells) or col_idx >= len(cells[row_idx]):
+                logger.error(f"Cell ({row_idx},{col_idx}) out of bounds")
+                return False
+
+            cell = cells[row_idx][col_idx]
+            insertion_index = cell.get("insertion_index")
+
+            if not insertion_index:
+                logger.warning(f"No insertion_index for cell ({row_idx},{col_idx})")
+                return False
+
+            # Insert text
+            await asyncio.to_thread(
+                self.service.documents()
+                .batchUpdate(
+                    documentId=document_id,
+                    body={
+                        "requests": [
+                            {
+                                "insertText": {
+                                    "location": {"index": insertion_index},
+                                    "text": cell_text,
+                                }
+                            }
+                        ]
+                    },
+                )
+                .execute
+            )
+
+            # Apply bold formatting if requested
+            if apply_bold:
+                await self._apply_bold_formatting(
+                    document_id, insertion_index, insertion_index + len(cell_text)
+                )
+
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to populate single cell: {str(e)}")
+            return False
+
+    async def _apply_bold_formatting(
+        self, document_id: str, start_index: int, end_index: int
+    ) -> None:
+        """Apply bold formatting to a text range."""
+        await asyncio.to_thread(
+            self.service.documents()
+            .batchUpdate(
+                documentId=document_id,
+                body={
+                    "requests": [
+                        {
+                            "updateTextStyle": {
+                                "range": {
+                                    "startIndex": start_index,
+                                    "endIndex": end_index,
+                                },
+                                "textStyle": {"bold": True},
+                                "fields": "bold",
+                            }
+                        }
+                    ]
+                },
+            )
+            .execute
+        )
+
+    async def populate_existing_table(
+        self,
+        document_id: str,
+        table_index: int,
+        table_data: List[List[str]],
+        clear_existing: bool = False,
+    ) -> Tuple[bool, str, Dict[str, Any]]:
+        """
+        Populate an existing table with data.
+
+        Args:
+            document_id: ID of the document
+            table_index: Index of the table to populate (0-based)
+            table_data: 2D list of data to insert
+            clear_existing: Whether to clear existing content first
+
+        Returns:
+            Tuple of (success, message, metadata)
+        """
+        try:
+            tables = await self._get_document_tables(document_id)
+            if table_index >= len(tables):
+                return (
+                    False,
+                    f"Table index {table_index} not found. Document has {len(tables)} tables",
+                    {},
+                )
+
+            table_info = tables[table_index]
+
+            # Validate dimensions
+            table_rows = table_info["rows"]
+            table_cols = table_info["columns"]
+            data_rows = len(table_data)
+            data_cols = len(table_data[0]) if table_data else 0
+
+            if data_rows > table_rows or data_cols > table_cols:
+                return (
+                    False,
+                    f"Data ({data_rows}x{data_cols}) exceeds table dimensions ({table_rows}x{table_cols})",
+                    {},
+                )
+
+            # Populate cells
+            population_count = await self._populate_existing_table_cells(
+                document_id, table_index, table_data
+            )
+
+            metadata = {
+                "table_index": table_index,
+                "populated_cells": population_count,
+                "table_dimensions": f"{table_rows}x{table_cols}",
+                "data_dimensions": f"{data_rows}x{data_cols}",
+            }
+
+            return (
+                True,
+                f"Successfully populated {population_count} cells in existing table",
+                metadata,
+            )
+
+        except Exception as e:
+            return False, f"Failed to populate existing table: {str(e)}", {}
+
+    async def _populate_existing_table_cells(
+        self, document_id: str, table_index: int, table_data: List[List[str]]
+    ) -> int:
+        """Populate cells in an existing table."""
+        population_count = 0
+
+        for row_idx, row_data in enumerate(table_data):
+            for col_idx, cell_text in enumerate(row_data):
+                if not cell_text:
+                    continue
+
+                # Get fresh table structure for each cell
+                tables = await self._get_document_tables(document_id)
+                if table_index >= len(tables):
+                    break
+
+                table = tables[table_index]
+                cells = table.get("cells", [])
+
+                if row_idx >= len(cells) or col_idx >= len(cells[row_idx]):
+                    continue
+
+                cell = cells[row_idx][col_idx]
+
+                # For existing tables, append to existing content
+                cell_end = cell["end_index"] - 1  # Don't include cell end marker
+
+                try:
+                    await asyncio.to_thread(
+                        self.service.documents()
+                        .batchUpdate(
+                            documentId=document_id,
+                            body={
+                                "requests": [
+                                    {
+                                        "insertText": {
+                                            "location": {"index": cell_end},
+                                            "text": cell_text,
+                                        }
+                                    }
+                                ]
+                            },
+                        )
+                        .execute
+                    )
+                    population_count += 1
+
+                except Exception as e:
+                    logger.error(
+                        f"Failed to populate existing cell ({row_idx},{col_idx}): {str(e)}"
+                    )
+
+        return population_count