""" OMML equation extraction from DOCX files. This module provides the OMMLExtractor class which extracts OMML (Office Math Markup Language) equations from Microsoft Word DOCX files, including main document, headers, and footers. """ import zipfile import logging from pathlib import Path from typing import List, Dict, Any from lxml import etree as ET logger = logging.getLogger(__name__) class OMMLExtractor: """Extract OMML equations from DOCX files. Handles extraction from: - Main document body (word/document.xml) - Headers (word/header*.xml) - Footers (word/footer*.xml) Preserves: - Non-breaking spaces (NBSP, U+00A0) - Display mode detection (inline vs display equations) - Context from surrounding paragraphs - Position tracking Example: >>> extractor = OMMLExtractor() >>> equations = extractor.extract_from_docx('document.docx') >>> print(f"Found {len(equations)} equations") >>> print(extractor.stats) """ # OMML and Word namespace definitions NAMESPACES = { 'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main', 'm': 'http://schemas.openxmlformats.org/officeDocument/2006/math', 'wp': 'http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing' } def __init__(self): """Initialize extractor with empty statistics.""" self.stats = { 'equations_found': 0, 'files_processed': 0, 'errors': 0 } def extract_from_docx(self, docx_path: str) -> List[Dict[str, Any]]: """Extract all OMML equations from a DOCX file. Args: docx_path: Path to DOCX file Returns: List of equation dictionaries with keys: - omml: Raw OMML XML string - cleaned_omml: Cleaned OMML ready for transformation - position: Position/index in document - context: Surrounding text context - source_file: XML file within DOCX (e.g., 'word/document.xml') - display_mode: Boolean indicating display (True) vs inline (False) Raises: FileNotFoundError: If DOCX file doesn't exist Exception: For other processing errors Example: >>> extractor = OMMLExtractor() >>> equations = extractor.extract_from_docx('test.docx') >>> for eq in equations: ... print(f"Position {eq['position']}: {eq['display_mode']}") """ equations = [] docx_path = Path(docx_path) if not docx_path.exists(): raise FileNotFoundError(f"DOCX file not found: {docx_path}") try: with zipfile.ZipFile(docx_path, 'r') as docx: # Extract main document equations.extend(self._extract_from_document(docx, 'word/document.xml')) # Check for equations in headers/footers header_files = [f for f in docx.namelist() if f.startswith('word/header')] footer_files = [f for f in docx.namelist() if f.startswith('word/footer')] for header_file in header_files: equations.extend(self._extract_from_document(docx, header_file)) for footer_file in footer_files: equations.extend(self._extract_from_document(docx, footer_file)) self.stats['files_processed'] += 1 self.stats['equations_found'] += len(equations) except Exception as e: self.stats['errors'] += 1 logger.error(f"Error processing {docx_path}: {e}") raise return equations def _extract_from_document(self, docx: zipfile.ZipFile, xml_path: str) -> List[Dict[str, Any]]: """Extract equations from a specific XML document within the DOCX. Args: docx: Open zipfile.ZipFile object xml_path: Path to XML file within DOCX Returns: List of equation dictionaries """ equations = [] try: if xml_path not in docx.namelist(): return equations doc_xml = docx.read(xml_path) # IMPORTANT: Preserve NBSP characters before XML parsing # lxml normalizes U+00A0 (NBSP) to regular space, so we replace it # with a placeholder that we'll restore later # UTF-8 encoding of NBSP is \xc2\xa0 doc_xml = doc_xml.replace(b'\xc2\xa0', b' ') # Convert NBSP to decimal entity # Use lxml for parsing to enable getparent() method root = ET.fromstring(doc_xml) # Find all OMML equations math_elements = root.findall('.//m:oMath', self.NAMESPACES) for i, math_elem in enumerate(math_elements): try: # Extract raw OMML (use lxml for consistency) raw_omml = ET.tostring(math_elem, encoding='unicode', method='xml') # Clean and prepare OMML cleaned_omml = self._clean_omml(raw_omml) # Get surrounding context context = self._get_context(math_elem) # Detect if this is a display equation (wrapped in oMathPara) display_mode = self._is_display_equation(math_elem) equation_data = { 'omml': raw_omml, 'cleaned_omml': cleaned_omml, 'position': i, 'context': context, 'source_file': xml_path, 'display_mode': display_mode } equations.append(equation_data) except Exception as e: logger.warning(f"Error extracting equation {i} from {xml_path}: {e}") continue except Exception as e: logger.error(f"Error reading {xml_path}: {e}") return equations def _clean_omml(self, omml: str) -> str: """Clean and prepare OMML for transformation. Fixes: - Missing namespace declarations - Word formatting properties - Whitespace normalization between tags Args: omml: Raw OMML XML string Returns: Cleaned OMML string """ import re # Add proper namespace declarations if missing if 'xmlns:m=' not in omml: omml = omml.replace( ']*?/>', '', omml) # Keep w:color elements for color preservation # omml = re.sub(r']*?/>', '', omml) # Normalize whitespace BETWEEN XML tags only (not inside text content) # Collapse multiple whitespace chars between tags into single space # Pattern: ">[ \t\n\r]+<" → "> <" omml = re.sub(r'>([ \t\n\r]+)<', r'> <', omml) omml = omml.strip() return omml def _get_context(self, math_elem: ET._Element) -> str: """Get surrounding text context for the equation. Args: math_elem: OMML math element Returns: Surrounding paragraph text (max 200 characters) """ try: # Find the parent paragraph current = math_elem while current is not None: if current.tag.endswith('p'): # Word paragraph break parent = current.find('..') current = parent if current is None: return "" # Extract all text from the paragraph text_parts = [] for elem in current.iter(): if elem.tag.endswith('t') and elem.text: # Word text element text_parts.append(elem.text) context = ' '.join(text_parts).strip() # Limit context length if len(context) > 200: context = context[:200] + "..." return context except Exception as e: logger.debug(f"Error getting context: {e}") return "" def _is_display_equation(self, math_elem: ET._Element) -> bool: """Determine if an equation is display mode ($$) or inline mode ($). Display equations are wrapped in , inline equations are not. Args: math_elem: OMML math element Returns: True for display mode, False for inline mode """ try: # Check if the parent is oMathPara parent = math_elem.getparent() if parent is not None and parent.tag.endswith('oMathPara'): return True return False except Exception as e: logger.debug(f"Error detecting display mode: {e}") return False # Default to inline if unsure