diff --git a/jsondoc/utils/block.py b/jsondoc/utils/block.py index 7155048..235c6aa 100644 --- a/jsondoc/utils/block.py +++ b/jsondoc/utils/block.py @@ -54,3 +54,70 @@ def _process_block_and_children( if hasattr(block, "children") and block.children: for child in block.children: _process_block_and_children(child, block_map) + + +def replace_urls( + input_obj: Page | BlockBase | list[BlockBase], + url_replace_map: dict[str, str], +) -> None: + """ + Recursively processes all blocks in the input object and replaces any file or image URLs + that are found in the url_replace_map with their corresponding values. + + Args: + input_obj: Can be either a Page object, a single Block object, or a list of Block objects + url_replace_map: A dictionary mapping original URLs to their replacement URLs + """ + # Extract all blocks + blocks = extract_blocks(input_obj) + + # Process each block + for block_id, block in blocks.items(): + _replace_url_in_block(block, url_replace_map) + + +def _replace_url_in_container(container, container_type, url_replace_map): + """ + Helper function to replace URL in a file or external container. + + Args: + container: The container object that might have URL information + container_type: The type of container ('file' or 'external') + url_replace_map: Dictionary mapping original URLs to replacement URLs + """ + if not hasattr(container, container_type): + return + + url_container = getattr(container, container_type) + if hasattr(url_container, "url"): + current_url = getattr(url_container, "url") + if current_url in url_replace_map: + setattr(url_container, "url", url_replace_map[current_url]) + + +def _replace_url_in_block(block: BlockBase, url_replace_map: dict[str, str]) -> None: + """ + Helper function to replace URLs in a specific block based on the url_replace_map. + + Args: + block: The block to process + url_replace_map: A dictionary mapping original URLs to their replacement URLs + """ + if not hasattr(block, "type"): + return + + block_type = getattr(block, "type") + + # Check for image block + if block_type == "image" and hasattr(block, "image"): + image = getattr(block, "image") + if hasattr(image, "type"): + image_type = getattr(image, "type") + _replace_url_in_container(image, image_type, url_replace_map) + + # Check for file block + elif block_type == "file" and hasattr(block, "file"): + file_block = getattr(block, "file") + if hasattr(file_block, "type"): + file_type = getattr(file_block, "type") + _replace_url_in_container(file_block, file_type, url_replace_map)