1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
|
# pyright: reportUnknownMemberType=false
# pyright: reportAttributeAccessIssue=false
# pyright: reportOptionalMemberAccess=false
from __future__ import annotations as _annotations
import os
import sys
from pathlib import Path
from typing import TYPE_CHECKING, cast
from bs4 import Tag
from typing_extensions import TypedDict
from pydantic import TypeAdapter
if TYPE_CHECKING:
from mkdocs.config import Config
from mkdocs.structure.files import Files
from mkdocs.structure.pages import Page
class AlgoliaRecord(TypedDict):
content: str
pageID: str
abs_url: str
title: str
objectID: str
rank: int
records: list[AlgoliaRecord] = []
records_ta = TypeAdapter(list[AlgoliaRecord])
# these values should match docs/javascripts/search-worker.js.
ALGOLIA_APP_ID = 'KPPUDTIAVX'
ALGOLIA_INDEX_NAME = 'pydantic-docs'
# Algolia has a limit of 100kb per record in the paid plan,
# leave some space for the other fields as well.
MAX_CONTENT_LENGTH = 90_000
def get_heading_text(heading: Tag):
return heading.get_text().replace('ΒΆ', '').strip().replace('\n', ' ')
def on_page_content(html: str, page: Page, config: Config, files: Files) -> str:
if not os.getenv('CI'):
return html
from bs4 import BeautifulSoup
assert page.title is not None, 'Page title must not be None'
title = cast(str, page.title)
soup = BeautifulSoup(html, 'html.parser')
# If the page does not start with a heading, add the h1 with the title
# Some examples don't have a heading. or start with h2
first_element = soup.find()
if (
not first_element
or not first_element.name # type: ignore[reportAttributeAccessIssue]
or first_element.name not in ['h1', 'h2', 'h3'] # type: ignore[reportAttributeAccessIssue]
):
soup.insert(0, BeautifulSoup(f'<h1 id="{title}">{title}</h1>', 'html.parser'))
# Clean up presentational and UI elements
for element in soup.find_all(['autoref']):
element.decompose()
# this removes the large source code embeds from Github
for element in soup.find_all('details'):
element.decompose()
# Cleanup code examples
for extra in soup.find_all('div', attrs={'class': 'language-python highlight'}):
extra.replace_with(BeautifulSoup(f'<pre>{extra.find("code").get_text()}</pre>', 'html.parser'))
# Cleanup code examples, part 2
for extra in soup.find_all('div', attrs={'class': 'language-python doc-signature highlight'}):
extra.replace_with(BeautifulSoup(f'<pre>{extra.find("code").get_text()}</pre>', 'html.parser'))
# The API reference generates HTML tables with line numbers, this strips the line numbers cell and goes back to a code block
for extra in soup.find_all('table', attrs={'class': 'highlighttable'}):
extra.replace_with(BeautifulSoup(f'<pre>{extra.find("code").get_text()}</pre>', 'html.parser'))
headings = soup.find_all(['h1', 'h2', 'h3'])
# Use the rank to put the sections in the beginning higher in the search results
rank = 100
# Process each section
for current_heading in headings:
heading_id = current_heading.get('id', '')
section_title = get_heading_text(current_heading) # type: ignore[reportArgumentType]
# Get content until next heading
content: list[str] = []
sibling = current_heading.find_next_sibling()
while sibling and sibling.name not in {'h1', 'h2', 'h3'}:
content.append(str(sibling))
sibling = sibling.find_next_sibling()
section_html = ''.join(content)
section_soup = BeautifulSoup(section_html, 'html.parser')
section_plain_text = section_soup.get_text(' ', strip=True)
# Create anchor URL
anchor_url: str = f'{page.abs_url}#{heading_id}' if heading_id else page.abs_url or ''
record_title = title
if current_heading.name == 'h2':
record_title = f'{title} - {section_title}'
elif current_heading.name == 'h3':
previous_heading: Tag = current_heading.find_previous(['h1', 'h2']) # type: ignore[reportAssignmentType]
record_title = f'{title} - {get_heading_text(previous_heading)} - {section_title}'
# print(f'Adding record {record_title}')
# Create record for this section
records.append(
AlgoliaRecord(
content=section_plain_text,
pageID=title,
abs_url=anchor_url,
title=record_title,
objectID=anchor_url,
rank=rank,
)
)
rank -= 5
return html
ALGOLIA_RECORDS_FILE = 'algolia_records.json'
def on_post_build(config: Config) -> None:
if records:
algolia_records_path = Path(config['site_dir']) / ALGOLIA_RECORDS_FILE
with algolia_records_path.open('wb') as f:
f.write(records_ta.dump_json(records))
def algolia_upload() -> None:
from algoliasearch.search.client import SearchClientSync
algolia_write_api_key = os.environ['ALGOLIA_WRITE_API_KEY']
client = SearchClientSync(ALGOLIA_APP_ID, algolia_write_api_key)
filtered_records: list[AlgoliaRecord] = []
algolia_records_path = Path.cwd() / 'site' / ALGOLIA_RECORDS_FILE
with algolia_records_path.open('rb') as f:
all_records = records_ta.validate_json(f.read())
for record in all_records:
content = record['content']
if len(content) > MAX_CONTENT_LENGTH:
print(
f"Record with title '{record['title']}' has more than {MAX_CONTENT_LENGTH} characters, {len(content)}."
)
print(content)
else:
filtered_records.append(record)
print(f'Uploading {len(filtered_records)} out of {len(all_records)} records to Algolia...')
client.clear_objects(index_name=ALGOLIA_INDEX_NAME)
client.set_settings(
index_name=ALGOLIA_INDEX_NAME,
index_settings={
'searchableAttributes': ['title', 'content'],
'attributesToSnippet': ['content:40'],
'customRanking': [
'desc(rank)',
],
},
)
client.batch(
index_name=ALGOLIA_INDEX_NAME,
batch_write_params={'requests': [{'action': 'addObject', 'body': record} for record in filtered_records]},
)
if __name__ == '__main__':
if sys.argv[-1] == 'upload':
algolia_upload()
else:
print('Run with "upload" argument to upload records to Algolia.')
sys.exit(1)
|
