Notebook

def translate_config(config: dict, data: pd.DataFrame) -> str:
    """Convert an OWID chart configuration into Python code.

    Takes a chart configuration dictionary (as returned by get_chart_config) and the
    corresponding data, and generates Python code that recreates the chart using the
    owid-grapher-py Chart API.

    Args:
        config: OWID chart configuration dictionary containing chart type, dimensions,
            selections, labels, and other settings.
        data: pandas DataFrame containing the chart data in long format (typically
            with 'year'/'date', 'entity', and 'value' columns).

    Returns:
        String containing Python code that recreates the chart. The code uses method
        chaining with Chart().mark_*().encode().label() etc.

    Raises:
        UnsupportedChartType: If the chart type is not yet supported for translation.
            Currently supports: LineChart.

    Example:
        ```python
        config = get_chart_config(slug='life-expectancy')
        data = get_chart_data(slug='life-expectancy')
        code = translate_config(config, data)
        print(code)
        # Output:
        # grapher.Chart(
        #     data
        # ).encode(
        #     x="year",
        #     y="value",
        #     entity="entity"
        # ).label(
        #     title="Life Expectancy"
        # )
        ```
    """
    # jsonschema.validate(config, WHITELIST_SCHEMA)

    chart_type = config.get("type", "LineChart")
    if chart_type == "LineChart":
        return translate_line_chart(config, data)

    raise UnsupportedChartType(chart_type)

def translate_line_chart(config: dict, data: pd.DataFrame) -> str:
    """Translate a LineChart configuration to Python code.

    Generates Python code for creating a line chart by analyzing the configuration
    and determining appropriate encoding, selection, labels, and interaction settings.

    Args:
        config: Chart configuration dictionary for a LineChart.
        data: DataFrame containing the chart data.

    Returns:
        Python code string for creating the line chart.
    """
    encoding = _gen_encoding(config, data)
    preselection, selection = _gen_selection(config, data)
    labels = _gen_labels(config)
    interaction = _gen_interaction(config)
    transform = _gen_transform(config)
    mark_map = _gen_mark_map(config)

    return f"""
grapher.Chart(
    data{preselection}
).mark_line(){mark_map}{encoding}{selection}{transform}{labels}{interaction}
""".strip()

def generate_notebook(config: dict, path: str) -> None:
    """Generate a Jupyter notebook from a chart configuration.

    Fetches the chart data, translates the config to Python code, and creates
    a complete Jupyter notebook with imports, data loading, and chart code.

    Args:
        config: OWID chart configuration dictionary (must include 'slug').
        path: Directory path where the notebook file will be saved. The notebook
            will be named {slug}.ipynb.

    Raises:
        UnsupportedChartType: If the chart type cannot be translated.
    """
    owid_data = get_owid_data(config)
    data = owid_data_to_frame(owid_data)
    py = translate_config(config, data)

    slug = config["slug"]
    title = config["title"]
    save_to_notebook(slug, title, py, path)

def save_to_notebook(slug: str, title: str, py: str, path: str) -> None:
    """Save chart code to a Jupyter notebook file.

    Creates a new notebook with the chart title, imports, data loading, and
    the generated Python code.

    Args:
        slug: Chart slug for the filename and data loading.
        title: Chart title for the notebook heading.
        py: Python code string to include in the notebook.
        path: Directory path where the notebook will be saved.
    """
    nb_file = join(path, f"{slug}.ipynb")

    nb = _new_notebook(slug, title, py)

    with open(nb_file, "w") as ostream:
        nbf.write(nb, ostream)

@click.command()
@click.argument("input_file")
@click.argument("dest_path")
def main(input_file, dest_path):
    """
    Take a large list of configs in JSONL format, and attempt to render as many as we can to
    notebooks.
    """
    data_folder = join(dest_path, "_data")
    if not isdir(data_folder):
        mkdir(data_folder)

    i = 0
    for total, config in enumerate(iter_published(input_file), 1):
        slug = config["slug"]
        try:
            generate_notebook(config, dest_path)
            i += 1
            print(click.style(f"✓ [{i}/{total}] {slug}", fg="green"))
        except UnsupportedChartType:
            print(f"✗ [{i}/{total}] {slug}")

        except Exception as e:
            print(f"ERROR: {slug}")
            raise e

    print(f"Generated {i} notebooks successfully")