Order image per step plots (#143)

mattseddon · web-flow · commit 6d045d7db92c · 2023-09-14T16:54:59.000-04:00
diff --git a/src/dvc_render/html.py b/src/dvc_render/html.py
@@ -1,6 +1,9 @@
+import os
 from pathlib import Path
 from typing import TYPE_CHECKING, List, Optional
 
+from dvc_render.image import ImageRenderer
+
 from .exceptions import DvcRenderException
 
 if TYPE_CHECKING:
@@ -76,6 +79,25 @@ def embed(self) -> str:
         return self.template
 
 
+def _order_image_per_step(renderer: "Renderer") -> tuple:
+    is_image_renderer = isinstance(renderer, ImageRenderer)
+
+    if not is_image_renderer:
+        return (is_image_renderer, None, None, None)
+
+    path = renderer.name
+    basename = os.path.basename(path)
+    filename = os.path.splitext(basename)[0]
+    image_number = int(filename) if filename.isdigit() else None
+
+    return (
+        is_image_renderer,
+        os.path.dirname(path),
+        image_number,
+        basename,
+    )
+
+
 def render_html(
     renderers: List["Renderer"],
     output_file: "StrPath",
@@ -94,7 +116,12 @@ def render_html(
 
     document = HTML(page_html, refresh_seconds=refresh_seconds)
 
-    for renderer in renderers:
+    sorted_renderers = sorted(
+        renderers,
+        key=_order_image_per_step,
+    )
+
+    for renderer in sorted_renderers:
         document.with_scripts(renderer.SCRIPTS)
         document.with_element(renderer.generate_html(html_path=output_path))
 
diff --git a/tests/test_html.py b/tests/test_html.py
@@ -1,7 +1,17 @@
 # pylint: disable=missing-function-docstring, R0801
+import os
+
 import pytest
 
-from dvc_render.html import HTML, PAGE_HTML, MissingPlaceholderError, render_html
+from dvc_render.html import (
+    HTML,
+    PAGE_HTML,
+    MissingPlaceholderError,
+    _order_image_per_step,
+    render_html,
+)
+from dvc_render.image import ImageRenderer
+from dvc_render.vega import VegaRenderer
 
 CUSTOM_PAGE_HTML = """<!DOCTYPE html>
 <html>
@@ -83,6 +93,45 @@ def test_render_html_with_custom_template(mocker, tmp_dir):
     assert output_file.read_text() == CUSTOM_PAGE_HTML.format(plot_divs="")
 
 
+def test_order_image_per_step():
+    image_per_step_dir = "dvclive"
+    other_image_dir = "static"
+
+    def create_renderer(filename: str) -> ImageRenderer:
+        return ImageRenderer(
+            [
+                {
+                    "filename": filename,
+                    "rev": "workspace",
+                    "src": filename,
+                }
+            ],
+            filename,
+        )
+
+    r1 = VegaRenderer([], "dvc.yaml::Loss")
+    r2 = VegaRenderer([], "dvc.yaml::Accuracy")
+    r3 = create_renderer(os.path.join(image_per_step_dir, "0.jpg"))
+    r4 = create_renderer(os.path.join(image_per_step_dir, "1.jpg"))
+    r5 = create_renderer(os.path.join(image_per_step_dir, "2.jpg"))
+    r6 = create_renderer(os.path.join(image_per_step_dir, "10.jpg"))
+    r7 = create_renderer(os.path.join(other_image_dir, "a_file.jpg"))
+    r8 = create_renderer(os.path.join(other_image_dir, "z_file.jpg"))
+
+    renderers = [r7, r3, r5, r8, r1, r6, r4, r2]
+
+    assert sorted(renderers, key=_order_image_per_step) == [
+        r1,
+        r2,
+        r3,
+        r4,
+        r5,
+        r6,
+        r7,
+        r8,
+    ]
+
+
 def test_no_placeholder():
     template = "<head></head><body></body>"
 
