1# Copyright (c) 2023 Valentin Goldite. All Rights Reserved.
  2"""Routines to manipulate the tags on the keys of a dict.
  3
  4Used by the processing objects.
  5"""
  6
  7import copy
  8import re
  9from typing import Any, Dict, List, Tuple
 10
 11
 12def clean_tag(flat_key: str, tag_name: str) -> str:
 13    """Clean a tag from a flat key.
 14
 15    It removes all occurrences of the tag with the exact name.
 16
 17    Parameters
 18    ----------
 19    flat_key : str
 20        The flat key to clean.
 21    tag_name : str
 22        The name of the tag to remove, with or without the '@' prefix.
 23
 24    Returns
 25    -------
 26    flat_key : str
 27        The cleaned flat key.
 28
 29    Examples
 30    --------
 31    >>> clean_tag('abc@tag.def@tag_2.ghi@tag', 'tag')
 32    abc.def@tag_2.ghi
 33
 34    .. note::
 35        `tag_name` is supposed to be the exact name of the tag.
 36    """
 37    if tag_name[0] == "@":
 38        tag_name = tag_name[1:]
 39    # Replace "@tag@other_tag" by "@other_tag"
 40    parts = flat_key.split(f"@{tag_name}@")
 41    flat_key = "@".join(parts)
 42    # Replace "@tag." by "."
 43    parts = flat_key.split(f"@{tag_name}.")
 44    flat_key = ".".join(parts)
 45    # Remove "@tag" at the end of the string
 46    if flat_key.endswith(f"@{tag_name}"):
 47        flat_key = flat_key[: -len(f"@{tag_name}")]
 48    return flat_key
 49
 50
 51def clean_all_tags(flat_key: str) -> str:
 52    """Clean all tags from a flat key.
 53
 54    Parameters
 55    ----------
 56    flat_key : str
 57        The flat key to clean.
 58
 59    Returns
 60    -------
 61    flat_key : str
 62        The cleaned flat key.
 63    """
 64    list_keys = flat_key.split(".")
 65    for i, key in enumerate(list_keys):
 66        key = re.sub(r"@.*", "", key)
 67        list_keys[i] = key
 68    flat_key = ".".join(list_keys)
 69    return flat_key
 70
 71
 72def dict_clean_tags(flat_dict: Dict[str, Any]) -> Tuple[Dict[str, Any], List[str]]:
 73    """Clean a dict from all tags and return the list of keys with tags.
 74
 75    Parameters
 76    ----------
 77    flat_dict : Dict[str, Any]
 78        The flat dict to clean.
 79
 80    Returns
 81    -------
 82    clean_dict : Dict[str, Any]
 83        The cleaned flat dict without tags in the keys.
 84    tagged_keys : List[str]
 85        The list of keys with tags that have been cleaned.
 86    """
 87    items = list(flat_dict.items())
 88    clean_dict = copy.deepcopy(flat_dict)
 89    tagged_keys = []
 90    for key, value in items:
 91        if "@" in key:
 92            del clean_dict[key]
 93            clean_dict[clean_all_tags(key)] = value
 94            tagged_keys.append(key)
 95    return clean_dict, tagged_keys
 96
 97
 98def is_tag_in(flat_key: str, tag_name: str, *, full_key: bool = False) -> bool:
 99    """Check if a tag is in a flat key.
100
101    The tag name must be the exact name, with or without the "@".
102    It supports the case where there are other tags that are prefixes
103    or suffixes of the considered tag.
104
105    Parameters
106    ----------
107    flat_key : str
108        The flat key to check.
109    tag_name : str
110        The name of the tag to check, with or without the '@' prefix.
111    full_key : bool, optional
112        If True, check for the full key. If False, check for the last part of
113        the flat key (after the last dot) that correspond to the parameter name.
114        By default, False.
115
116    Returns
117    -------
118    bool
119        True if the tag is in the flat key, False otherwise.
120    """
121    if tag_name[0] == "@":
122        tag_name = tag_name[1:]
123    if not full_key:
124        flat_key = flat_key.split(".")[-1]
125    is_in = (
126        flat_key.endswith(f"@{tag_name}")
127        or f"@{tag_name}@" in flat_key
128        or f"@{tag_name}." in flat_key
129    )
130    return is_in