dvc/tests/integration/plots/test_plots.py

import json
import os
from copy import deepcopy
from typing import Dict, List
from urllib.parse import urlparse
from urllib.request import url2pathname

import dpath.util
import pytest
from bs4 import BeautifulSoup
from funcy import first

from dvc.cli import main
from dvc.render import REVISION_FIELD, VERSION_FIELD

JSON_OUT = "vis_data"


def call(capsys, subcommand="show"):
    capsys.readouterr()
    assert (
        main(["plots", subcommand, "--json", "-o", JSON_OUT, "--split"]) == 0
    )
    split_json_out, _ = capsys.readouterr()

    split_json_result = json.loads(split_json_out)

    capsys.readouterr()
    assert main(["plots", subcommand, "--json", "-o", JSON_OUT]) == 0
    json_out, _ = capsys.readouterr()

    json_result = json.loads(json_out)

    assert main(["plots", subcommand]) == 0
    html_path_out, _ = capsys.readouterr()

    parsed = urlparse(html_path_out.strip())
    abspath = url2pathname(parsed.path)
    return abspath, json_result, split_json_result


def extract_vega_specs(html_path, plots_ids):
    from dvc_render.base import Renderer

    result = {}

    with open(html_path, encoding="utf-8") as fd:
        content = fd.read()

    reader = BeautifulSoup(content, features="html.parser")
    for plot_id in plots_ids:
        script = _remove_blanks(
            reader.find(
                "div", id=Renderer.remove_special_chars(plot_id)
            ).script.text
        )
        result[plot_id] = json.loads(
            script.split("; vegaEmbed")[0].replace("var spec = ", "")
        )

    return result


def drop_fields(datapoints: List[Dict], fields: List[str]):
    tmp = deepcopy(datapoints)
    for datapoint in tmp:
        keys = set(datapoint.keys())
        for key in keys:
            if key in fields:
                datapoint.pop(key)
    return tmp


def verify_image(path, version, filename, content, html_path, json_result):
    assert os.path.exists(html_path)
    with open(html_path, encoding="utf-8") as fd:
        html_content = fd.read()

    image_data = {}
    for datapoint in json_result[filename]:
        if datapoint["revisions"] == [version]:
            image_data = datapoint
            break

    assert image_data, f"{version} data for {filename} was not found"
    assert image_data["type"] == "image"
    output_filename = filename.replace("/", "_")
    output_name = f"{version}_{output_filename}"
    assert image_data["url"] == str(path / JSON_OUT / output_name)
    assert (path / JSON_OUT / output_name).read_bytes() == content

    assert os.path.join("static", output_name) in html_content

    # there should be no absolute paths in produced HTML
    # TODO uncomment once dvc-render is adjusted
    # assert str(path) not in html_content
    assert (
        path / "dvc_plots" / "static" / output_name
    ).read_bytes() == content


def _remove_blanks(text: str):
    return " ".join(text.replace("\t", "").replace("\n", "").split())


class hashabledict(dict):
    def __hash__(self):
        return hash(tuple(sorted(self.items())))


def verify_vega(
    versions,
    html_result,
    json_result,
    split_json_result,
):
    if isinstance(versions, str):
        versions = [versions]

    for j in [json_result, split_json_result]:
        assert len(j) == 1
        assert j[0]["type"] == "vega"
        assert set(j[0]["revisions"]) == set(versions)

    assert json_result[0]["datapoints"] == split_json_result[0]["datapoints"]
    assert set(versions) == set(json_result[0]["datapoints"].keys())

    assert json_result[0]["content"]["data"]["values"]
    assert html_result["data"]["values"]
    assert (
        split_json_result[0]["content"]["data"]["values"]
        == "<DVC_METRIC_DATA>"
    )

    def _assert_templates_equal(
        html_template, filled_template, split_template
    ):
        # besides data, json and split json should be equal
        path = ["data", "values"]
        tmp1 = deepcopy(html_template)
        tmp2 = deepcopy(filled_template)
        tmp3 = deepcopy(split_template)
        dpath.util.set(tmp1, path, {})
        dpath.util.set(tmp2, path, {})
        dpath.util.set(tmp3, path, {})

        assert tmp1 == tmp2 == tmp3

    _assert_templates_equal(
        html_result, json_result[0]["content"], split_json_result[0]["content"]
    )


def verify_vega_props(plot_id, json_result, title, x, y, **kwargs):
    data = json_result[plot_id]
    assert len(data) == 1
    data = first(data)

    assert dpath.util.get(data, ["content", "title"]) == title

    try:
        # TODO confusion_matrix_plot - need to find better way of asserting
        #      encoding as its place is not constant in vega
        plot_x = dpath.util.get(
            data, ["content", "spec", "encoding", "x", "field"]
        )
        plot_y = dpath.util.get(
            data, ["content", "spec", "encoding", "y", "field"]
        )
    except KeyError:
        # default plot
        plot_x = dpath.util.get(
            data, ["content", "layer", 0, "encoding", "x", "field"]
        )
        plot_y = dpath.util.get(
            data, ["content", "layer", 0, "encoding", "y", "field"]
        )

    assert plot_x == x
    assert plot_y == y


def _update_datapoints(datapoints: List, update: Dict):
    result = []
    for dp in datapoints:
        tmp = dp.copy()
        tmp.update(update)
        result.append(tmp)
    return result


@pytest.mark.vscode
def test_no_plots(tmp_dir, scm, dvc, capsys):
    html_path, json_result, split_json_result = call(capsys)
    assert not os.path.exists(html_path)
    assert json_result == {}
    assert split_json_result == {}


@pytest.mark.vscode
def test_repo_with_plots(
    tmp_dir, scm, dvc, capsys, run_copy_metrics, repo_with_plots
):
    repo_state = repo_with_plots()

    image_v1, linear_v1, confusion_v1, confusion_props = next(repo_state)

    html_path, json_result, split_json_result = call(capsys)
    html_result = extract_vega_specs(
        html_path, ["linear.json", "confusion.json"]
    )

    assert json_result["linear.json"][0]["content"]["data"][
        "values"
    ] == _update_datapoints(
        linear_v1,
        {
            VERSION_FIELD: {
                "revision": "workspace",
                "filename": "linear.json",
                "field": "y",
            },
        },
    )
    assert html_result["linear.json"]["data"]["values"] == _update_datapoints(
        linear_v1,
        {
            REVISION_FIELD: "workspace",
        },
    )
    assert json_result["confusion.json"][0]["content"]["data"][
        "values"
    ] == _update_datapoints(
        confusion_v1,
        {
            VERSION_FIELD: {
                "revision": "workspace",
                "filename": "confusion.json",
                "field": "actual",
            },
        },
    )
    assert html_result["confusion.json"]["data"][
        "values"
    ] == _update_datapoints(
        confusion_v1,
        {
            REVISION_FIELD: "workspace",
        },
    )
    verify_image(
        tmp_dir, "workspace", "image.png", image_v1, html_path, json_result
    )

    for plot in ["linear.json", "confusion.json"]:
        verify_vega(
            "workspace",
            html_result[plot],
            json_result[plot],
            split_json_result[plot],
        )

    verify_vega_props("confusion.json", json_result, **confusion_props)

    image_v2, linear_v2, confusion_v2, confusion_props = next(repo_state)

    html_path, json_result, split_json_result = call(capsys, subcommand="diff")
    html_result = extract_vega_specs(
        html_path, ["linear.json", "confusion.json"]
    )

    verify_image(
        tmp_dir, "workspace", "image.png", image_v2, html_path, json_result
    )
    verify_image(
        tmp_dir, "HEAD", "image.png", image_v1, html_path, json_result
    )

    for plot in ["linear.json", "confusion.json"]:
        verify_vega(
            ["HEAD", "workspace"],
            html_result[plot],
            json_result[plot],
            split_json_result[plot],
        )
    verify_vega_props("confusion.json", json_result, **confusion_props)
    path = tmp_dir / "subdir"
    path.mkdir()
    with path.chdir():
        html_path, json_result, split_json_result = call(
            capsys, subcommand="diff"
        )
        html_result = extract_vega_specs(
            html_path,
            ["../linear.json", "../confusion.json"],
        )
        assert json_result["../linear.json"][0]["content"]["data"][
            "values"
        ] == _update_datapoints(
            linear_v2,
            {
                VERSION_FIELD: {
                    "revision": "workspace",
                    "filename": "../linear.json",
                    "field": "y",
                },
            },
        ) + _update_datapoints(
            linear_v1,
            {
                VERSION_FIELD: {
                    "revision": "HEAD",
                    "filename": "../linear.json",
                    "field": "y",
                },
            },
        )
        assert html_result["../linear.json"]["data"][
            "values"
        ] == _update_datapoints(
            linear_v2,
            {
                REVISION_FIELD: "workspace",
            },
        ) + _update_datapoints(
            linear_v1,
            {
                REVISION_FIELD: "HEAD",
            },
        )
        assert json_result["../confusion.json"][0]["content"]["data"][
            "values"
        ] == _update_datapoints(
            confusion_v2,
            {
                VERSION_FIELD: {
                    "revision": "workspace",
                    "filename": "../confusion.json",
                    "field": "actual",
                },
            },
        ) + _update_datapoints(
            confusion_v1,
            {
                VERSION_FIELD: {
                    "revision": "HEAD",
                    "filename": "../confusion.json",
                    "field": "actual",
                },
            },
        )
        assert html_result["../confusion.json"]["data"][
            "values"
        ] == _update_datapoints(
            confusion_v2,
            {
                REVISION_FIELD: "workspace",
            },
        ) + _update_datapoints(
            confusion_v1,
            {
                REVISION_FIELD: "HEAD",
            },
        )

        for plot in [
            "../linear.json",
            "../confusion.json",
        ]:
            verify_vega(
                ["HEAD", "workspace"],
                html_result[plot],
                json_result[plot],
                split_json_result[plot],
            )
        verify_image(
            path,
            "workspace",
            "../image.png",
            image_v2,
            html_path,
            json_result,
        )
        verify_image(
            path,
            "HEAD",
            "../image.png",
            image_v1,
            html_path,
            json_result,
        )


@pytest.mark.vscode
def test_repo_with_removed_plots(tmp_dir, capsys, repo_with_plots):
    from dvc.utils.fs import remove

    next(repo_with_plots())

    # even if there is no data, call should be successful
    remove(tmp_dir / ".dvc" / "cache")
    remove("linear.json")
    remove("confusion.json")
    remove("image.png")

    for s in {"show", "diff"}:
        _, json_result, split_json_result = call(capsys, subcommand=s)
        for p in {
            "linear.json",
            "confusion.json",
            "image.png",
        }:
            assert json_result[p] == []
            assert split_json_result[p] == []


def test_config_output_dir(tmp_dir, dvc, capsys):
    subdir = tmp_dir / "subdir"
    ret = main(["config", "plots.out_dir", os.fspath(subdir)])
    assert ret == 0

    metric = [{"first_val": 100, "val": 2}, {"first_val": 200, "val": 3}]
    (tmp_dir / "metric.json").dump_json(metric, sort_keys=True)

    assert main(["plots", "show", "metric.json"]) == 0

    out, _ = capsys.readouterr()
    assert subdir.as_uri() in out
    assert subdir.is_dir()
    assert (subdir / "index.html").is_file()

    cli_arg_subdir = tmp_dir / "cli_option"
    assert (
        main(["plots", "show", "-o", os.fspath(cli_arg_subdir), "metric.json"])
        == 0
    )

    out, _ = capsys.readouterr()
    assert cli_arg_subdir.as_uri() in out
    assert cli_arg_subdir.is_dir()
    assert (cli_arg_subdir / "index.html").is_file()


@pytest.mark.vscode
def test_repo_with_config_plots(tmp_dir, capsys, repo_with_config_plots):
    repo_state = repo_with_config_plots()
    plots = next(repo_state)

    html_path, _, __ = call(capsys)

    assert os.path.exists(html_path)
    html_result = extract_vega_specs(
        html_path,
        [
            "dvc.yaml::linear_train_vs_test",
            "dvc.yaml::confusion_train_vs_test",
        ],
    )
    ble = _update_datapoints(
        plots["data"]["linear_train.json"],
        {
            REVISION_FIELD: "linear_train.json",
        },
    ) + _update_datapoints(
        plots["data"]["linear_test.json"],
        {
            REVISION_FIELD: "linear_test.json",
        },
    )

    assert (
        html_result["dvc.yaml::linear_train_vs_test"]["data"]["values"] == ble
    )
    # TODO check json results once vscode is able to handle flexible plots