conan-io/conan

import os

import yaml

from conans.errors import ConanException
from conans.util.files import load, save


def update_conandata(conanfile, data):
    """
    this only works for updating the conandata on the export() method, it seems it would
    be plain wrong to try to change it anywhere else
    """
    if not hasattr(conanfile, "export_folder") or conanfile.export_folder is None:
        raise ConanException("The 'update_conandata()' can only be used in the 'export()' method")
    path = os.path.join(conanfile.export_folder, "conandata.yml")
    if os.path.exists(path):
        conandata = load(path)
        conandata = yaml.safe_load(conandata)
    else:  # File doesn't exist, create it
        conandata = {}

    def recursive_dict_update(d, u):
        for k, v in u.items():
            if isinstance(v, dict):
                d[k] = recursive_dict_update(d.get(k, {}), v)
            else:
                d[k] = v
        return d

    recursive_dict_update(conandata, data)
    new_content = yaml.safe_dump(conandata)
    save(path, new_content)


def trim_conandata(conanfile):
    """
    Tool to modify the ``conandata.yml`` once it is exported, to limit it to the current version
    only
    """
    if not hasattr(conanfile, "export_folder") or conanfile.export_folder is None:
        raise ConanException("The 'trim_conandata()' can only be used in the 'export()' method")
    path = os.path.join(conanfile.export_folder, "conandata.yml")
    if not os.path.exists(path):
        raise ConanException("conandata.yml file doesn't exist")

    conandata = load(path)
    conandata = yaml.safe_load(conandata)

    version = str(conanfile.version)
    result = {}
    for k, v in conandata.items():
        if not isinstance(v, dict):
            result[k] = v
            continue  # to allow user extra conandata, common to all versions
        version_data = v.get(version)
        if version_data is not None:
            result[k] = {version: version_data}

    new_conandata_yml = yaml.safe_dump(result, default_flow_style=False)
    save(path, new_conandata_yml)