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