Some code to help with manifest parsing

[lswainemoore]

Hi--the parsing of the dist folder here:

flask-vite/src/flask_vite/tags.py

Line 25 in 25357ee

js_file = Path(glob.glob(f"{vite_folder_path}/dist/assets/*.js")[0]).name

(just taking the first found css/js files) wasn't cutting it for my purposes, because it was picking the wrong one, preventing my site from loading.

I saw that proper parsing of the manifest.json file that Vite can produce is a TODO (https://github.com/abilian/flask-vite/blame/25357ee0e4bde73c31b9b96526cbdee4f87d9eb7/TODO.md#L6), so I went ahead and implemented it myself. I think it would be great to get this feature into this library, so I'll include most of my code below--it's not quite shaped as a PR, but hopefully it makes it easier to tackle this.

A couple notes before the code:

• I've tried to follow the instructions from Vite given here: https://vite.dev/guide/backend-integration.
  • This means I implemented their pseudocode for the manifest parsing, and used their examples as unit tests.
  • It also means that I made a couple changes to the expected entrypoint. In step 1, Vite recommends replacing the index.html with your actual JS file. I did this as well, though I made the entrypoint configurable in the tag.
  • The ordering of the static files changes a bit in the header, and I'm including the modulepreload links.
  • I'm using react, so I did the extra bit in step 2 (including some extra code alongside debug tags) to fix hot reloading.
• Making this work requires changing one's vite.config.js to output the manifest. In case that's not desired, I'd imagine one might want to fallback to the current behavior. The code below does that, but in a slightly awkward way from the POV of integrating it into the library--that is, it uses the current version flask-vite's make_static_tag. That would need to change.
• There's a wrinkle around the assets prefix and the fact that that folder is baked into the static route--I didn't want to change the routing, so I handled it in a fairly ad hoc way. I'm not sure how exactly flask-vite would want to handle this.
• YMMV: this works for my purposes, but I can't vouch it's got all the edge cases covered.

custom_vite.py:

import json

from flask import current_app, url_for
from flask_vite import Vite, tags
from markupsafe import Markup


# flask-vite is lacking in a couple ways.
# we extend it to solve them in this file.
# ideally, most of this code gets merged to flask-vite at some point.


def imported_chunks(manifest, name):
    """
    Pulls relevant files from manifest.json.
    Adapted from pseudocode at https://vite.dev/guide/backend-integration.html.
    For more on why we're doing this, see better_vite_tags.
    """
    seen = set()

    def get_imported_chunks(chunk):
        chunks = []
        for file in chunk.get("imports", []):
            importee = manifest[file]
            if file in seen:
                continue
            seen.add(file)
            chunks.extend(get_imported_chunks(importee))
            chunks.append(importee)
        return chunks

    return get_imported_chunks(manifest[name])


def links_from_manifest(manifest, name):
    """
    Order links that vite manifest implies we should use.
    Also adapted from: https://vite.dev/guide/backend-integration.html
    For more on why we're doing this, see better_vite_tags.
    """
    links = []

    for css_file in manifest[name].get("css", []):
        links.append(("stylesheet", css_file))

    chunks = imported_chunks(manifest, name)

    for chunk in chunks:
        for css_file in chunk.get("css", []):
            links.append(("stylesheet", css_file))

    links.append(("module", manifest[name]["file"]))

    for chunk in chunks:
        links.append(("modulepreload", chunk["file"]))

    return links


def better_make_debug_tag():
    return """
        <!-- FLASK_VITE_HEADER (dev) -->
        
        <!-- helps with the hot reloading. from: https://vite.dev/guide/backend-integration.html -->
        <script type="module">
            import RefreshRuntime from 'http://localhost:3000/@react-refresh'
            RefreshRuntime.injectIntoGlobalHook(window)
            window.$RefreshReg$ = () => {}
            window.$RefreshSig$ = () => (type) => type
            window.__vite_plugin_react_preamble_installed__ = true
        </script>

        <script type="module" src="http://localhost:3000/@vite/client"></script>
        <script type="module" src="http://localhost:3000/main.js"></script>
    """


def better_make_static_tag(entrypoint):
    this_ext = current_app.extensions["vite"]
    vite_folder_path = this_ext.vite_folder_path
    manifest = this_ext.manifest

    # fallback to the library's version if we don't have a manifest
    if manifest is None:
        return tags.make_static_tag()

    links = links_from_manifest(manifest, entrypoint)

    def make_tag(tag_type, file):
        # this is a bit subtle, but flask-vite wants us to provide only the filename itself
        # and it'll add /assets/ on its own
        # (https://github.com/abilian/flask-vite/blob/25357ee0e4bde73c31b9b96526cbdee4f87d9eb7/src/flask_vite/extension.py#L121)
        # (at least as of 0.5.2).
        # but the manifest by default references files with "assets" included in the filename.
        if file.startswith("assets/"):
            file = file.replace("assets/", "")
        url = url_for(f"{vite_folder_path}.static", filename=file)
        if tag_type == "stylesheet":
            return f'<link rel="stylesheet" href="{url}" />'
        if tag_type == "module":
            return f'<script type="module" src="{url}"></script>'
        if tag_type == "modulepreload":
            return f'<link rel="modulepreload" href="{url}" />'

    tags_str = "\n".join([make_tag(tag_type, file) for tag_type, file in links])

    return f"""
        <!-- FLASK_VITE_HEADER (build) -->
        {tags_str}
    """


def better_make_tag(*, static=False, entrypoint="main.js"):
    # this is more or less what happens in flask-vite's make_tag
    # except that that's got a couple problems:
    # 1) for development, it's not configured to work well with vite-react (and hot reloading)
    # 2) for production, they pick random js/css file from dist folder with code like this:
    #       js_file = Path(glob.glob(f"{vite_folder_path}/dist/assets/*.js")[0]).name
    #       css_file = Path(glob.glob(f"{vite_folder_path}/dist/assets/*.css")[0]).name
    if static or not current_app.debug:
        tag = better_make_static_tag(entrypoint)
    else:
        tag = better_make_debug_tag()
    return Markup(tag)


class CustomVite(Vite):
    def init_app(self, app, *args, **kwargs):
        super().init_app(app, *args, **kwargs)

        # read this only once.
        try:
            with open(f"{self.vite_folder_path}/dist/.vite/manifest.json") as f:
                self.manifest = json.load(f)
        except FileNotFoundError:
            # this is primarily for CI environments, but we'll
            # fail gracefully when this file isn't there.
            self.manifest = None

        app.template_global("better_vite_tags")(better_make_tag)

test_custom_vite.py:

import unittest

from app.custom_vite import links_from_manifest


class ViteTestCase(unittest.TestCase):
    # just test our manifest.json stuff for now--no need for
    # db or anything.
    def test_links_from_manifest(self):
        # this is from: https://vite.dev/guide/backend-integration.html
        manifest = {
            "_shared-B7PI925R.js": {
                "file": "assets/shared-B7PI925R.js",
                "name": "shared",
                "css": ["assets/shared-ChJ_j-JJ.css"],
            },
            "_shared-ChJ_j-JJ.css": {
                "file": "assets/shared-ChJ_j-JJ.css",
                "src": "_shared-ChJ_j-JJ.css",
            },
            "baz.js": {
                "file": "assets/baz-B2H3sXNv.js",
                "name": "baz",
                "src": "baz.js",
                "isDynamicEntry": True,
            },
            "views/bar.js": {
                "file": "assets/bar-gkvgaI9m.js",
                "name": "bar",
                "src": "views/bar.js",
                "isEntry": True,
                "imports": ["_shared-B7PI925R.js"],
                "dynamicImports": ["baz.js"],
            },
            "views/foo.js": {
                "file": "assets/foo-BRBmoGS9.js",
                "name": "foo",
                "src": "views/foo.js",
                "isEntry": True,
                "imports": ["_shared-B7PI925R.js"],
                "css": ["assets/foo-5UjPuW-k.css"],
            },
        }

        # these are the two examples from:
        # https://vite.dev/guide/backend-integration.html
        self.assertEqual(
            links_from_manifest(manifest, "views/foo.js"),
            [
                ("stylesheet", "assets/foo-5UjPuW-k.css"),
                ("stylesheet", "assets/shared-ChJ_j-JJ.css"),
                ("module", "assets/foo-BRBmoGS9.js"),
                ("modulepreload", "assets/shared-B7PI925R.js"),
            ],
        )
        self.assertEqual(
            links_from_manifest(manifest, "views/bar.js"),
            [
                ("stylesheet", "assets/shared-ChJ_j-JJ.css"),
                ("module", "assets/bar-gkvgaI9m.js"),
                ("modulepreload", "assets/shared-B7PI925R.js"),
            ],
        )

I then use CustomVite just as I would the current library Vite, when init-ing my app.

Let me know if you have any questions about this--I'd be happy to help with getting changes like these merged upstream, so that I don't have to maintain this code myself.