ygrep

#!/usr/bin/env python3

import os
import yaml
import argparse
import fnmatch

# CC-Zero, this code was generated by AI
# Flatten a bunch of yaml dicts and lists, and grep for a matching key


def flatten_yaml(data, parent_key='', result=None, list_key='name'):
    """
    Recursively flatten a nested YAML structure into key-value pairs.
    If lists are encountered, replace indices with a value from list_key if available.

    Args:
        data (any): The current level of the YAML data.
        parent_key (str): The key path leading to this level.
        result (dict): The flattened result dictionary.
        list_key (str): The key to use as the name for list items.

    Returns:
        dict: A dictionary of flattened key-value pairs.
    """
    if result is None:
        result = {}

    if isinstance(data, dict):
        for key, value in data.items():
            new_key = f"{parent_key}.{key}" if parent_key else key
            flatten_yaml(value, new_key, result, list_key)
    elif isinstance(data, list):
        for item in data:
            if isinstance(item, dict) and list_key in item:
                item_name = item[list_key]
                new_key = f"{parent_key}.['{item_name}']"
            else:
                item_name = str(data.index(item))
                new_key = f"{parent_key}.[{item_name}]"
            flatten_yaml(item, new_key, result, list_key)
    else:
        result[parent_key] = data

    return result


def search_yaml_files(search_key, paths='.', list_key='name'):
    """
    Search for a specific key pattern in all YAML files under a directory.

    Args:
        search_key (str): The key path pattern to search for.
        paths (list): The directories to search in.
        list_key (str): The key to use for naming list items.
    """
    for base_dir in paths:
      for root, _, files in os.walk(base_dir):
          for file in files:
              if file.endswith(('.yaml', '.yml')):
                  file_path = os.path.join(root, file)
                  try:
                      with open(file_path, "r") as f:
                          data = yaml.safe_load(f)
                          if data:
                              flattened = flatten_yaml(data, list_key=list_key)
                              for key, value in flattened.items():
                                  if fnmatch.fnmatch(key, search_key):
                                      print(f"{file_path}: {key} = {value}")
                  except Exception as e:
                      print(f"Error reading {file_path}: {e}")


def main():
    parser = argparse.ArgumentParser(
        description="Search YAML files for a specific flattened key pattern.")
    parser.add_argument("search_key", help="The key path pattern to search for (e.g., 'some.dict.[*].with.lists').")
    parser.add_argument("paths", nargs="*", default=".", help="The directories to search in.")
    parser.add_argument("--list-key", default="name", help="The key used to name list items (default: 'name').")
    args = parser.parse_args()

    search_yaml_files(args.search_key, args.paths, args.list_key)


if __name__ == "__main__":
    main()