capnpc-py/src/capnp_stubgen/utils.py

"""Utility functions for capnp-stubgen."""

import os
from typing import Any

#: A dict from capnp's ``CodeGeneratorRequest.to_dict()``.
#: Keys are strings, values are arbitrary nested capnp data.
NodeDict = dict[str, Any]


def filename_to_module_name(filename: str) -> str:
    """Convert a .capnp filename to the corresponding Python module name.

    Example: ``path/to/my_file.capnp`` → ``my_file_capnp``
    """
    base = os.path.basename(filename)
    if base.endswith(".capnp"):
        base = base[:-6]
    return base + "_capnp"


def get_display_name(node: NodeDict) -> str:
    """Extract the short display name from a node dict."""
    return str(node["displayName"])[int(node["displayNamePrefixLength"]):]


def get_node_kind(node: NodeDict) -> str:
    """Determine the kind of a node from its dict representation.

    Returns one of: "file", "struct", "enum", "interface", "const", "annotation".
    """
    for kind in ("file", "struct", "enum", "interface", "const", "annotation"):
        if kind in node:
            return kind
    return "unknown"


def type_id_to_int(type_id: str) -> int:
    """Convert a type ID string (from dict) to an integer."""
    return int(type_id)