221 lines
8.4 KiB
Python
221 lines
8.4 KiB
Python
"""
|
||
File writer module for sphinx-llms-txt.
|
||
"""
|
||
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Tuple, Union
|
||
|
||
from sphinx.application import Sphinx
|
||
from sphinx.util import logging
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def _fix_mojibake(text: str) -> str:
|
||
"""Fix common UTF-8/Windows-1252 mojibake in text.
|
||
|
||
Mojibake (文字化け, "character transformation") is a Japanese term for garbled text
|
||
caused by decoding bytes with the wrong character encoding.
|
||
|
||
This handles the case where UTF-8 bytes were incorrectly decoded as Windows-1252
|
||
(or Latin-1), resulting in corrupted characters like:
|
||
- ' (U+2019) becoming ’
|
||
- " (U+201C) becoming “
|
||
- " (U+201D) becoming â€
|
||
- — (U+2014) becoming â€"
|
||
- – (U+2013) becoming â€"
|
||
- … (U+2026) becoming …
|
||
|
||
Args:
|
||
text: The potentially corrupted text string
|
||
|
||
Returns:
|
||
The repaired text string, or the original if no repair was needed/possible
|
||
"""
|
||
if not text:
|
||
return text
|
||
|
||
try:
|
||
# Try to encode the text as Windows-1252 (which would succeed if it contains
|
||
# the mojibake characters) and then decode as UTF-8 (to get the original)
|
||
return text.encode("windows-1252").decode("utf-8")
|
||
except (UnicodeDecodeError, UnicodeEncodeError):
|
||
# If encoding/decoding fails, the text is either:
|
||
# - Already correct UTF-8
|
||
# - Corrupted in a different way we can't fix
|
||
return text
|
||
|
||
|
||
class FileWriter:
|
||
"""Handles writing processed content to output files."""
|
||
|
||
def __init__(self, config: Dict[str, Any], outdir: str = None, app: Sphinx = None):
|
||
self.config = config
|
||
self.outdir = outdir
|
||
self.app = app
|
||
|
||
def _resolve_uri_template(self, sources_dir: Path = None) -> str:
|
||
"""Resolve which URI template to use based on configuration and sources_dir.
|
||
|
||
Args:
|
||
sources_dir: Path to _sources directory (None if not found)
|
||
|
||
Returns:
|
||
The template string to use for generating URIs
|
||
"""
|
||
# If custom template exists
|
||
custom_template = self.config.get("llms_txt_uri_template")
|
||
|
||
if custom_template:
|
||
# Validate user's template by checking for valid variable names
|
||
try:
|
||
# Try formatting with test valid values to validate syntax
|
||
test_values = {
|
||
"base_url": "http://example.com/",
|
||
"docname": "test",
|
||
"suffix": ".rst",
|
||
"sourcelink_suffix": ".txt",
|
||
}
|
||
custom_template.format(**test_values)
|
||
return custom_template
|
||
except (KeyError, ValueError) as e:
|
||
logger.warning(
|
||
f"sphinx-llms-txt: Invalid llms_txt_uri_template: {e}. "
|
||
f"Falling back to default."
|
||
)
|
||
|
||
# Else, use one of the default templates
|
||
if sources_dir:
|
||
return "{base_url}_sources/{docname}{suffix}{sourcelink_suffix}"
|
||
else:
|
||
return "{base_url}{docname}.html"
|
||
|
||
def write_combined_file(
|
||
self, content_parts: List[str], output_path: Path, total_line_count: int
|
||
) -> bool:
|
||
"""Write the combined content to a file.
|
||
|
||
Args:
|
||
content_parts: List of content strings to combine
|
||
output_path: Path to write the output file
|
||
total_line_count: Total number of lines in the content
|
||
|
||
Returns:
|
||
True if successful, False otherwise
|
||
"""
|
||
try:
|
||
with open(output_path, "w", encoding="utf-8") as f:
|
||
f.write("\n".join(content_parts))
|
||
|
||
logger.info(
|
||
f"sphinx-llms-txt: Created {output_path} with {len(content_parts)}"
|
||
f" sources and {total_line_count} lines"
|
||
)
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"sphinx-llms-txt: Error writing combined sources file: {e}")
|
||
return False
|
||
|
||
def write_verbose_info_to_file(
|
||
self,
|
||
page_order: Union[List[str], List[Tuple[str, str]]],
|
||
page_titles: Dict[str, str],
|
||
total_line_count: int = 0,
|
||
sources_dir: Path = None,
|
||
) -> bool:
|
||
"""Write summary information to the llms.txt file.
|
||
|
||
Args:
|
||
page_order: Ordered list of document names or (docname, suffix) tuples
|
||
page_titles: Dictionary mapping docnames to titles
|
||
total_line_count: Total number of lines in the combined content
|
||
sources_dir: Path to _sources directory (None if not found)
|
||
|
||
Returns:
|
||
True if successful, False otherwise
|
||
"""
|
||
if not self.outdir:
|
||
logger.warning(
|
||
"sphinx-llms-txt: Cannot write verbose info to file: outdir not set"
|
||
)
|
||
return False
|
||
|
||
output_path = Path(self.outdir) / self.config.get("llms_txt_filename")
|
||
try:
|
||
with open(output_path, "w", encoding="utf-8") as f:
|
||
project_name = "llms-txt Summary"
|
||
# First priority: use title from config if available
|
||
if self.config.get("llms_txt_title"):
|
||
project_name = self.config.get("llms_txt_title")
|
||
# Second priority: use project name from Sphinx app if available
|
||
elif (
|
||
self.app
|
||
and hasattr(self.app, "config")
|
||
and hasattr(self.app.config, "project")
|
||
):
|
||
project_name = self.app.config.project
|
||
# Fix any UTF-8/Windows-1252 mojibake in the project name
|
||
project_name = _fix_mojibake(project_name)
|
||
f.write(f"# {project_name}\n\n")
|
||
|
||
# Add description if available
|
||
description = self.config.get("llms_txt_summary", "")
|
||
if description:
|
||
# Trim leading and trailing whitespace
|
||
description = description.strip()
|
||
if description:
|
||
# Fix any UTF-8/Windows-1252 mojibake in the description
|
||
description = _fix_mojibake(description)
|
||
# Only add blockquote if description is not empty
|
||
# Replace newlines with newline + blockquote marker to maintain
|
||
# blockquote formatting
|
||
description = description.replace("\n", "\n> ")
|
||
f.write(f"> {description}\n\n")
|
||
|
||
f.write("## Docs\n\n")
|
||
# Get base URL from config
|
||
base_url = self.config.get("html_baseurl", "/")
|
||
# Ensure base_url ends with a trailing slash
|
||
if not base_url.endswith("/"):
|
||
base_url += "/"
|
||
|
||
# Get sourcelink suffix from Sphinx config
|
||
sourcelink_suffix = ""
|
||
if self.app and hasattr(self.app.config, "html_sourcelink_suffix"):
|
||
sourcelink_suffix = self.app.config.html_sourcelink_suffix
|
||
# Handle empty string case specially
|
||
if sourcelink_suffix == "":
|
||
sourcelink_suffix = "" # Keep it empty
|
||
elif not sourcelink_suffix.startswith("."):
|
||
sourcelink_suffix = "." + sourcelink_suffix
|
||
|
||
# Resolve which template to use
|
||
uri_template = self._resolve_uri_template(sources_dir)
|
||
|
||
for item in page_order:
|
||
# Handle both old format (str) and new format (tuple)
|
||
if isinstance(item, tuple):
|
||
docname, suffix = item
|
||
else:
|
||
docname = item
|
||
suffix = None
|
||
|
||
title = page_titles.get(docname, docname)
|
||
# Fix any UTF-8/Windows-1252 mojibake in the title
|
||
title = _fix_mojibake(title)
|
||
|
||
uri = uri_template.format(
|
||
base_url=base_url,
|
||
docname=docname,
|
||
suffix=suffix or "",
|
||
sourcelink_suffix=sourcelink_suffix,
|
||
)
|
||
|
||
f.write(f"- [{title}]({uri})\n")
|
||
|
||
logger.info(f"sphinx-llms-txt: created {output_path}")
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"sphinx-llms-txt: Error writing verbose info to file: {e}")
|
||
return False
|