MkDocs hook to add block override validation

validate_block.py

"""MkDocs hook to add block override validation

This hook adds a filter function to the Environment `validate_block`.
This function is a context filter, and takes the content, name, depth and saved checksum.

```jinja
{% block site_nav %} 
  {% set _ = super() | validate_block("site_nav", 2, "68f1518853") %}
  {{ super() }}
{% endblock %}

{% block site_nav %} 
  {% set _ = super() | validate_block("site_nav", 2) %}
  {{ super() }}
{% endblock %}
```
The `super()` is the content.
The `site_nav` is the name of the block, and cache name for hash
The `2` is the depth to limit the scope that is targeted for hash calculation
The `68f1518853` is the checksum used to compare against calculated checksum.
Without providing the checksum it will print out the calculated checksum.

If the checksum doesn't match, a warning is being logged, so --strict flag can help block 
deployments, but it can also be used as reference for local builds.

I got the idea after helping out with:

- https://github.com/squidfunk/mkdocs-material/discussions/8094

MIT License 2025 Kamil Krzyśków (HRY)
"""

import hashlib
import logging

from jinja2 import Environment
from mkdocs.plugins import PrefixedLogger
from mkdocs.structure.files import Files
from mkdocs.utils.templates import TemplateContext, contextfilter


def on_config(*_, **__):
    # Clear for serve reload
    ValidationData.reset()


def on_files(files: Files, *_, **__):
    if files.get_file_from_path(ValidationData.path_to_check) is None:
        LOG.warning(
            f"Path to validate the override blocks can't be found, '{ValidationData.path_to_check}'"
        )


def on_env(env: Environment, *_, **__):
    env.filters["validate_block"] = validate_block


def on_post_build(*_, **__):
    if not ValidationData.has_been_checked:
        LOG.warning(
            f"Template override validation has never run, maybe the hook should be removed from mkdocs.ymls?"
        )


@contextfilter
def validate_block(
    context: TemplateContext, content: str, name: str, depth: int, checksum: str = ""
) -> None:
    page = context.get("page")

    # Skip 404 static templates
    if page is None:
        return

    file_uri = page.file.src_uri

    if file_uri != ValidationData.path_to_check:
        return

    assert depth > 0, f"Depth must be greater than 0, {file_uri} -> {name}"
    checksum = checksum.strip()

    # Load checksum from cache
    calculated_checksum = ValidationData.cache.get(name)

    if calculated_checksum is None:
        # Truncate content for stable hashing
        content = truncate_html_by_depth(content, depth, file_uri)
        calculated_checksum: str = hashlib.sha256(content.encode(encoding="utf-8")).hexdigest()[:10]

    LOG.debug(">\n".join(content.split(">")))
    log_func = LOG.info

    if checksum and calculated_checksum != checksum:
        log_func = LOG.warning
        log_func(f"Checksum for block '{name}' in '{file_uri}' doesn't match provided '{checksum}'")

    if checksum == "" or calculated_checksum != checksum:
        log_func(
            f"Checksum for block '{name}' == '{calculated_checksum}', pass it in the function to validate later"
        )

    ValidationData.cache[name] = calculated_checksum
    ValidationData.has_been_checked = True


def truncate_html_by_depth(content: str, depth: int, file_uri: str) -> str:
    """Extract the tags in the tree limited by depth"""

    # Normalize whitespace
    content = " ".join(content.split()).replace("> ", ">").strip()
    stack = []
    result = []

    i = content.find("<", 0)
    if i < 0:
        LOG.warning(f"Could not find start tag in '{file_uri}'")
        return ""

    n = len(content)
    current_depth = len(stack)
    found_tags = False

    while i < n:
        if content[i] != "<":
            i += 1
            continue

        if content[i + 1] == "/":
            # Handle closing tag
            j = content.find(">", i)
            k = content.find("<", i + 1)  # Check if there is no opener sooner
            assert j > -1 and (k < 0 or j < k), "Closing tag not properly formed"
            close_pos = j + 1
            tag = content[i:close_pos]
            i = close_pos

            # Add close tag to output
            if current_depth <= depth:
                result.append(tag)

            # Assert stack is valid
            tag_name = tag.split("/")[1].rstrip(">").strip()
            assert tag_name == stack[-1], f"Stack invalid, expected {tag_name}, got {stack[-1]}"
            stack.pop()
            current_depth = len(stack)

            continue

        # Handle opening tag
        j = content.find(">", i)
        k = content.find("<", i + 1)  # Check if there is no opener sooner
        assert j > -1 and (k < 0 or j < k), "Opening tag not properly closed"
        close_pos = j + 1
        tag_body = content[i:close_pos]
        i = close_pos

        # Extract tag name (e.g., "div" from "<div class='xxx'>")
        tag_name = tag_body.split()[0].lstrip("<").rstrip(">").strip()
        is_self_closing = tag_name in [
            "br",
            "img",
            "input",
            "path",
        ]  # Add more self-closing tags as needed

        # Self closing tags don't add depth
        if not is_self_closing:
            stack.append(tag_name)
            current_depth = len(stack)

        # Add tag body to output
        if current_depth <= depth:
            result.append(tag_body)

        # Move to the next tag
        i = content.find("<", i)
        if i < 0:
            assert not found_tags or (
                found_tags and current_depth == 0
            ), "Loop finished without processing all tags"
            break

    return "".join(result)


HOOK_NAME: str = "validate_blocks"
"""Name of this hook. Used in logging."""


LOG: PrefixedLogger = PrefixedLogger(HOOK_NAME, logging.getLogger(f"mkdocs.hooks.{HOOK_NAME}"))
"""Logger instance for this hook."""


class ValidationData:
    cache: dict[str, str] = {}
    """Cache block_name hashes. For consistency always load it on the same page."""

    has_been_checked: bool = False
    """Flag to make sure the validation executed at least once"""

    path_to_check: str = "index.md"
    """Limit the hash validation for one path to improve performance"""

    @classmethod
    def reset(cls):
        cls.cache.clear()
        cls.has_been_checked = False