zulip/zerver/lib/markdown/static.py

37 lines
1.1 KiB
Python

from typing import Any
from xml.etree.ElementTree import Element
import markdown
from django.contrib.staticfiles.storage import staticfiles_storage
from markdown.extensions import Extension
from typing_extensions import override
from zerver.lib.markdown.priorities import PREPROCESSOR_PRIORITIES
class MarkdownStaticImagesGenerator(Extension):
@override
def extendMarkdown(self, md: markdown.Markdown) -> None:
md.treeprocessors.register(
StaticImageProcessor(md),
"static_images",
PREPROCESSOR_PRIORITIES["static_images"],
)
class StaticImageProcessor(markdown.treeprocessors.Treeprocessor):
"""
Rewrite img tags which refer to /static/ to use staticfiles
"""
@override
def run(self, root: Element) -> None:
for img in root.iter("img"):
url = img.get("src")
if url is not None and url.startswith("/static/"):
img.set("src", staticfiles_storage.url(url.removeprefix("/static/")))
def makeExtension(*args: Any, **kwargs: str) -> MarkdownStaticImagesGenerator:
return MarkdownStaticImagesGenerator(*args, **kwargs)