Helper function for `altair` exports

[timvink]

altair has good support for exporting vega schemas (https://altair-viz.github.io/user_guide/saving_charts.html?highlight=to_).

Of course you could just use chart.save('chart.svg') and insert that into a mkdocs page. But you want to want update the data later on and use the same chart.

Manually copy pasting the schema to a file is a hassle. Write a helper function to:

save_mkdocs_chart(chart, name="", docs_folder="../docs", json_path="docs/assets/charts/", data_path="docs/assets/charts/data")

The docs_folders could be auto-detected by traversing upward, finding mkdocs.yml and reading the docs_dir entry (if not present then docs/).

The json_path is relative to the docs_folder and can have a default "docs/assets/charts/"

The data_path is relative to the docs_folder and can have a "docs/assets/charts/data" default. This will only work if we can access the data object in the altair chart instance, which we can then write to .csv.

[ewellinger]

This might need to be tweaked somewhat, but I created this helper function that roughly follows the logic you're suggesting here:

import json
from pathlib import Path

from altair import Chart
from pydantic import BaseModel, DirectoryPath, validator


class ChartParams(BaseModel):
    output_dir: DirectoryPath
    data_file: str
    plot_name: str

    @validator("data_file")
    def validate_data_file(cls, val, values):
        if "output_dir" in values:
            full_path = values["output_dir"].absolute() / val
            if not full_path.exists():
                raise ValueError("data_file does not exist")
        return val

    @validator("plot_name")
    def validate_plot_name(cls, val):
        if not val.endswith("json"):
            raise ValueError(f"plot_name must end with .json")
        return val

    @property
    def plot_path(self):
        return self.output_dir / self.plot_name


def save_chart(chart: Chart, output_dir: Path, data_file: str = "results.csv", plot_name: str = "plot.json") -> None:
    params = ChartParams(output_dir=output_dir, data_file=data_file, plot_name=plot_name)

    # Convert the plot to JSON
    plot = json.loads(chart.to_json())

    del plot["datasets"]
    plot["data"] = {"url": data_file}
    with open(params.plot_path, "w") as fo:
        json.dump(plot, fo)

This assumes you're storing the plot in the same directory as the CSV file containing the data.