minimal python yaml query script

pyq.py

#!/usr/bin/env python3

import sys
import yaml
import os


def include_yaml(loader, node):
    # Get the path of the file to include
    file_path = os.path.join(os.path.dirname(loader.name), node.value)
    with open(file_path, "r") as file:
        return yaml.safe_load(file)


yaml.add_constructor("!include", include_yaml)


def get_yaml_section(file_path, key):
    try:
        # Read the YAML file
        with open(file_path, "r") as file:
            # Use a custom Loader to handle !include
            yaml_data = yaml.load(file, Loader=yaml.FullLoader)

        # Check if the key exists in the YAML data
        if key in yaml_data:
            return yaml_data[key]
        else:
            return f"Error: Key '{key}' not found in the YAML file."

    except FileNotFoundError:
        return f"Error: File '{file_path}' not found."
    except yaml.YAMLError as e:
        return f"Error: Invalid YAML format - {str(e)}"
    except Exception as e:
        return f"Error: An unexpected error occurred - {str(e)}"


if __name__ == "__main__":
    # Check if correct number of arguments are provided
    if len(sys.argv) != 3:
        print("Usage: pyq <file_path> <key>")
        sys.exit(1)

    file_path = sys.argv[1]
    key = sys.argv[2]

    # Get and print the specified section
    result = get_yaml_section(file_path, key)
    print(result)