Skip to content

Commit 2263163

Browse files
committed
perf: Small performance improvement when getting aliases and formatting/highlighting signatures
1 parent d754326 commit 2263163

4 files changed

Lines changed: 62 additions & 19 deletions

File tree

‎src/mkdocstrings_handlers/python/_internal/handler.py‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -365,23 +365,26 @@ def get_aliases(self, identifier: str) -> tuple[str, ...]:
365365
The aliases.
366366
"""
367367
if "(" in identifier:
368-
identifier, parameter = identifier.split("(", 1)
368+
identifier, _, parameter = identifier.partition("(")
369369
parameter = parameter.removesuffix(")")
370370
else:
371371
parameter = ""
372372
try:
373373
data = self._modules_collection[identifier]
374374
except (KeyError, AliasResolutionError):
375375
return ()
376-
aliases = [data.path]
376+
aliases: dict[str, object] = {data.path: None}
377377
try:
378-
for alias in [data.canonical_path, *data.aliases]:
379-
if alias not in aliases:
380-
aliases.append(alias)
378+
canonical_path = data.canonical_path
379+
data_aliases = data.aliases
381380
except AliasResolutionError:
382381
pass
382+
else:
383+
aliases[canonical_path] = None
384+
aliases.update(data_aliases)
383385
if parameter:
384-
return tuple(f"{alias}({parameter})" for alias in aliases)
386+
suffix = f"({parameter})"
387+
return tuple(alias + suffix for alias in aliases)
385388
return tuple(aliases)
386389

387390
def normalize_extension_paths(self, extensions: Sequence) -> list[str | dict[str, Any]]:

‎src/mkdocstrings_handlers/python/_internal/rendering.py‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,15 @@ def __call__(self, crossref: str, *, length: int) -> str:
125125
do_stash_crossref = _StashCrossRefFilter()
126126
"""Filter to stash cross-references (and restore them after formatting and highlighting)."""
127127

128+
_highlighting_span_re = re.compile(r'<span class="[a-z]{1,2}">')
129+
_crossref_key_re = re.compile(r"\b_[A-Za-z0-9]{2,}\b")
130+
131+
132+
def _restore_crossrefs(text: str, stash: dict[str, str]) -> str:
133+
text = _crossref_key_re.sub(lambda match: stash.get(match[0], match[0]), text)
134+
stash.clear()
135+
return text
136+
128137

129138
def _format_signature(name: Markup, signature: str, line_length: int) -> str:
130139
name = str(name).strip() # type: ignore[assignment]
@@ -199,13 +208,11 @@ def do_format_signature(
199208
# Pygments will set an `fm` (function -> magic) CSS class.
200209
# To fix this, we replace the CSS class in the first span with `nf`,
201210
# unless we already found an `nf` span.
202-
if not re.search(r'<span class="nf">', signature):
203-
signature = re.sub(r'<span class="[a-z]{1,2}">', '<span class="nf">', signature, count=1)
211+
if '<span class="nf">' not in signature:
212+
signature = _highlighting_span_re.sub('<span class="nf">', signature, count=1)
204213

205214
if stash := env.filters["stash_crossref"].stash:
206-
for key, value in stash.items():
207-
signature = re.sub(rf"\b{key}\b", value, signature)
208-
stash.clear()
215+
signature = _restore_crossrefs(signature, stash)
209216

210217
return signature
211218

@@ -267,9 +274,7 @@ def do_format_attribute(
267274
)
268275

269276
if stash := env.filters["stash_crossref"].stash:
270-
for key, value in stash.items():
271-
signature = re.sub(rf"\b{key}\b", value, signature)
272-
stash.clear()
277+
signature = _restore_crossrefs(signature, stash)
273278

274279
return signature
275280

@@ -322,13 +327,11 @@ def do_format_type_alias(
322327
# but instead as a regular name: `n` CSS class instead of `nc`.
323328
# To fix it, we replace the first occurrence of an `n` CSS class
324329
# with an `nc` one, unless we found `nc` already.
325-
if not re.search(r'<span class="nc">', signature):
326-
signature = re.sub(r'<span class="[a-z]{1,2}">', '<span class="nc">', signature, count=1)
330+
if '<span class="nc">' not in signature:
331+
signature = _highlighting_span_re.sub('<span class="nc">', signature, count=1)
327332

328333
if stash := env.filters["stash_crossref"].stash:
329-
for key, value in stash.items():
330-
signature = re.sub(rf"\b{key}\b", value, signature)
331-
stash.clear()
334+
signature = _restore_crossrefs(signature, stash)
332335

333336
return signature
334337

‎tests/test_handler.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import mkdocstrings
1616
import pytest
1717
from griffe import (
18+
Alias,
1819
Docstring,
1920
DocstringSectionExamples,
2021
DocstringSectionKind,
@@ -52,6 +53,32 @@ def test_collect_with_null_parser(handler: PythonHandler) -> None:
5253
assert handler.collect("mkdocstrings", PythonOptions(docstring_style=None))
5354

5455

56+
def test_get_aliases(handler: PythonHandler) -> None:
57+
"""Assert aliases are unique, ordered, and extended with parameter identifiers."""
58+
with temporary_visited_module(
59+
"def target(parameter): ...",
60+
module_name="package",
61+
modules_collection=handler._modules_collection,
62+
) as module:
63+
target = module["target"]
64+
module["alias"] = Alias("alias", target=target)
65+
module["other_alias"] = Alias("other_alias", target=target)
66+
module["unresolved"] = Alias("unresolved", target="missing.target")
67+
68+
assert handler.get_aliases("package.target") == (
69+
"package.target",
70+
"package.alias",
71+
"package.other_alias",
72+
)
73+
assert handler.get_aliases("package.alias(parameter)") == (
74+
"package.alias(parameter)",
75+
"package.target(parameter)",
76+
"package.other_alias(parameter)",
77+
)
78+
assert handler.get_aliases("package.unresolved") == ("package.unresolved",)
79+
assert handler.get_aliases("package.missing") == ()
80+
81+
5582
@pytest.mark.parametrize(
5683
"handler",
5784
[

‎tests/test_rendering.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,16 @@ def test_format_signature(name: Markup, signature: str) -> None:
5656
assert rendering._format_signature(name, signature, length)
5757

5858

59+
def test_restore_crossrefs() -> None:
60+
"""Restore multiple stashed cross-references without replacing partial keys."""
61+
stash = {"_short": "<autoref>short</autoref>", "_shorter": "<autoref>shorter</autoref>"}
62+
63+
restored = rendering._restore_crossrefs("_short + _shorter + x_short + _shorterx", stash)
64+
65+
assert restored == "<autoref>short</autoref> + <autoref>shorter</autoref> + x_short + _shorterx"
66+
assert not stash
67+
68+
5969
@dataclass
6070
class _FakeObject:
6171
name: str

0 commit comments

Comments
 (0)