"""
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,
) -> 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
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)
# Step 2: Get fresh document structure to find actual cell positions
fresh_tables = await self._get_document_tables(document_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
)
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
) -> 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)]},
)
.execute
)
async def _get_document_tables(self, document_id: str) -> List[Dict[str, Any]]:
"""Get fresh document structure and extract table information."""
doc = await asyncio.to_thread(
self.service.documents().get(documentId=document_id).execute
)
return find_tables(doc)
async def _populate_table_cells(
self, document_id: str, table_data: List[List[str]], bold_headers: bool
) -> 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,
)
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,
) -> 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)
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
