pydantic_graph.mermaid

DEFAULT_HIGHLIGHT_CSS module-attribute

DEFAULT_HIGHLIGHT_CSS = 'fill:#fdff32'

用于高亮节点的默认 CSS。

StateDiagramDirection module-attribute

StateDiagramDirection = Literal['TB', 'LR', 'RL', 'BT']

用于指定 mermaid 生成的状态图的方向。

• 'TB': 从上到下，这是 mermaid 图表的默认设置。
• 'LR': 从左到右
• 'RL': 从右到左
• 'BT': 从下到上

generate_code

generate_code(
    graph: Graph[Any, Any, Any],
    /,
    *,
    start_node: (
        Sequence[NodeIdent] | NodeIdent | None
    ) = None,
    highlighted_nodes: (
        Sequence[NodeIdent] | NodeIdent | None
    ) = None,
    highlight_css: str = DEFAULT_HIGHLIGHT_CSS,
    title: str | None = None,
    edge_labels: bool = True,
    notes: bool = True,
    direction: StateDiagramDirection | None,
) -> str

为图表生成 Mermaid 状态图 代码。

参数

名称	类型	描述	默认值
graph	Graph[Any, Any, Any]	要为其生成图像的图表。	必需
start_node	Sequence[NodeIdent] | NodeIdent | None	作为图表起点的节点标识符。	None
highlighted_nodes	Sequence[NodeIdent] | NodeIdent | None	要高亮的节点标识符。	None
highlight_css	str	用于高亮节点的 CSS。	DEFAULT_HIGHLIGHT_CSS
title	str | None	图表的标题。	None
edge_labels	bool	是否在图表中包含边标签。	True
notes	bool	是否在图表中包含注释。	True
direction	StateDiagramDirection | None	流动的方向。	必需

返回

类型	描述
str	图表的 Mermaid 代码。

源代码位于 pydantic_graph/pydantic_graph/mermaid.py

def generate_code(  # noqa: C901
    graph: Graph[Any, Any, Any],
    /,
    *,
    start_node: Sequence[NodeIdent] | NodeIdent | None = None,
    highlighted_nodes: Sequence[NodeIdent] | NodeIdent | None = None,
    highlight_css: str = DEFAULT_HIGHLIGHT_CSS,
    title: str | None = None,
    edge_labels: bool = True,
    notes: bool = True,
    direction: StateDiagramDirection | None,
) -> str:
    """Generate [Mermaid state diagram](https://mermaid.npmjs.net.cn/syntax/stateDiagram.html) code for a graph.

    Args:
        graph: The graph to generate the image for.
        start_node: Identifiers of nodes that start the graph.
        highlighted_nodes: Identifiers of nodes to highlight.
        highlight_css: CSS to use for highlighting nodes.
        title: The title of the diagram.
        edge_labels: Whether to include edge labels in the diagram.
        notes: Whether to include notes in the diagram.
        direction: The direction of flow.


    Returns:
        The Mermaid code for the graph.
    """
    start_node_ids = set(_node_ids(start_node or ()))
    for node_id in start_node_ids:
        if node_id not in graph.node_defs:
            raise LookupError(f'Start node "{node_id}" is not in the graph.')

    lines: list[str] = []
    if title:
        lines = ['---', f'title: {title}', '---']
    lines.append('stateDiagram-v2')
    if direction is not None:
        lines.append(f'  direction {direction}')
    for node_id, node_def in graph.node_defs.items():
        # we use round brackets (rounded box) for nodes other than the start and end
        if node_id in start_node_ids:
            lines.append(f'  [*] --> {node_id}')
        if node_def.returns_base_node:
            for next_node_id in graph.node_defs:
                lines.append(f'  {node_id} --> {next_node_id}')
        else:
            for next_node_id, edge in node_def.next_node_edges.items():
                line = f'  {node_id} --> {next_node_id}'
                if edge_labels and edge.label:
                    line += f': {edge.label}'
                lines.append(line)
        if end_edge := node_def.end_edge:
            line = f'  {node_id} --> [*]'
            if edge_labels and end_edge.label:
                line += f': {end_edge.label}'
            lines.append(line)

        if notes and node_def.note:
            lines.append(f'  note right of {node_id}')
            # mermaid doesn't like multiple paragraphs in a note, and shows if so
            clean_docs = re.sub('\n{2,}', '\n', node_def.note)
            lines.append(indent(clean_docs, '    '))
            lines.append('  end note')

    if highlighted_nodes:
        lines.append('')
        lines.append(f'classDef highlighted {highlight_css}')
        for node_id in _node_ids(highlighted_nodes):
            if node_id not in graph.node_defs:
                raise LookupError(f'Highlighted node "{node_id}" is not in the graph.')
            lines.append(f'class {node_id} highlighted')

    return '\n'.join(lines)

request_image

request_image(
    graph: Graph[Any, Any, Any],
    /,
    **kwargs: Unpack[MermaidConfig],
) -> bytes

使用 mermaid.ink 生成 Mermaid 图表的图像。

参数

名称	类型	描述	默认值
graph	Graph[Any, Any, Any]	要为其生成图像的图表。	必需
**kwargs	Unpack[MermaidConfig]	用于配置 mermaid 图表生成的其他参数。	{}

返回

类型	描述
bytes	图像数据。

源代码位于 pydantic_graph/pydantic_graph/mermaid.py

def request_image(
    graph: Graph[Any, Any, Any],
    /,
    **kwargs: Unpack[MermaidConfig],
) -> bytes:
    """Generate an image of a Mermaid diagram using [mermaid.ink](https://mermaid.ink).

    Args:
        graph: The graph to generate the image for.
        **kwargs: Additional parameters to configure mermaid chart generation.

    Returns:
        The image data.
    """
    code = generate_code(
        graph,
        start_node=kwargs.get('start_node'),
        highlighted_nodes=kwargs.get('highlighted_nodes'),
        highlight_css=kwargs.get('highlight_css', DEFAULT_HIGHLIGHT_CSS),
        title=kwargs.get('title'),
        edge_labels=kwargs.get('edge_labels', True),
        notes=kwargs.get('notes', True),
        direction=kwargs.get('direction'),
    )
    code_base64 = base64.b64encode(code.encode()).decode()

    params: dict[str, str | float] = {}
    if kwargs.get('image_type') == 'pdf':
        url = f'https://mermaid.ink/pdf/{code_base64}'
        if kwargs.get('pdf_fit'):
            params['fit'] = ''
        if kwargs.get('pdf_landscape'):
            params['landscape'] = ''
        if pdf_paper := kwargs.get('pdf_paper'):
            params['paper'] = pdf_paper
    elif kwargs.get('image_type') == 'svg':
        url = f'https://mermaid.ink/svg/{code_base64}'
    else:
        url = f'https://mermaid.ink/img/{code_base64}'

        if image_type := kwargs.get('image_type'):
            params['type'] = image_type

    if background_color := kwargs.get('background_color'):
        params['bgColor'] = background_color
    if theme := kwargs.get('theme'):
        params['theme'] = theme
    if width := kwargs.get('width'):
        params['width'] = width
    if height := kwargs.get('height'):
        params['height'] = height
    if scale := kwargs.get('scale'):
        params['scale'] = scale

    httpx_client = kwargs.get('httpx_client') or httpx.Client()
    response = httpx_client.get(url, params=params)
    if not response.is_success:
        raise httpx.HTTPStatusError(
            f'{response.status_code} error generating image:\n{response.text}',
            request=response.request,
            response=response,
        )
    return response.content

save_image

save_image(
    path: Path | str,
    graph: Graph[Any, Any, Any],
    /,
    **kwargs: Unpack[MermaidConfig],
) -> None

使用 mermaid.ink 生成 Mermaid 图表的图像并将其保存到本地文件。

参数

名称	类型	描述	默认值
path	Path | str	保存图像的路径。	必需
graph	Graph[Any, Any, Any]	要为其生成图像的图表。	必需
**kwargs	Unpack[MermaidConfig]	用于配置 mermaid 图表生成的其他参数。	{}

源代码位于 pydantic_graph/pydantic_graph/mermaid.py

def save_image(
    path: Path | str,
    graph: Graph[Any, Any, Any],
    /,
    **kwargs: Unpack[MermaidConfig],
) -> None:
    """Generate an image of a Mermaid diagram using [mermaid.ink](https://mermaid.ink) and save it to a local file.

    Args:
        path: The path to save the image to.
        graph: The graph to generate the image for.
        **kwargs: Additional parameters to configure mermaid chart generation.
    """
    if isinstance(path, str):
        path = Path(path)

    if 'image_type' not in kwargs:
        ext = path.suffix.lower()[1:]
        # no need to check for .jpeg/.jpg, as it is the default
        if ext in ('png', 'webp', 'svg', 'pdf'):
            kwargs['image_type'] = ext

    image_data = request_image(graph, **kwargs)
    path.write_bytes(image_data)

MermaidConfig

基类: TypedDict

用于配置 mermaid 图表生成的参数。

源代码位于 pydantic_graph/pydantic_graph/mermaid.py

class MermaidConfig(TypedDict, total=False):
    """Parameters to configure mermaid chart generation."""

    start_node: Sequence[NodeIdent] | NodeIdent
    """Identifiers of nodes that start the graph."""
    highlighted_nodes: Sequence[NodeIdent] | NodeIdent
    """Identifiers of nodes to highlight."""
    highlight_css: str
    """CSS to use for highlighting nodes."""
    title: str | None
    """The title of the diagram."""
    edge_labels: bool
    """Whether to include edge labels in the diagram."""
    notes: bool
    """Whether to include notes on nodes in the diagram, defaults to true."""
    image_type: Literal['jpeg', 'png', 'webp', 'svg', 'pdf']
    """The image type to generate. If unspecified, the default behavior is `'jpeg'`."""
    pdf_fit: bool
    """When using image_type='pdf', whether to fit the diagram to the PDF page."""
    pdf_landscape: bool
    """When using image_type='pdf', whether to use landscape orientation for the PDF.

    This has no effect if using `pdf_fit`.
    """
    pdf_paper: Literal['letter', 'legal', 'tabloid', 'ledger', 'a0', 'a1', 'a2', 'a3', 'a4', 'a5', 'a6']
    """When using image_type='pdf', the paper size of the PDF."""
    background_color: str
    """The background color of the diagram.

    If None, the default transparent background is used. The color value is interpreted as a hexadecimal color
    code by default (and should not have a leading '#'), but you can also use named colors by prefixing the
    value with `'!'`. For example, valid choices include `background_color='!white'` or `background_color='FF0000'`.
    """
    theme: Literal['default', 'neutral', 'dark', 'forest']
    """The theme of the diagram. Defaults to 'default'."""
    width: int
    """The width of the diagram."""
    height: int
    """The height of the diagram."""
    scale: Annotated[float, Ge(1), Le(3)]
    """The scale of the diagram.

    The scale must be a number between 1 and 3, and you can only set a scale if one or both of width and height are set.
    """
    httpx_client: httpx.Client
    """An HTTPX client to use for requests, mostly for testing purposes."""
    direction: StateDiagramDirection
    """The direction of the state diagram."""

start_node instance-attribute

start_node: Sequence[NodeIdent] | NodeIdent

作为图表起点的节点标识符。

highlighted_nodes instance-attribute

highlighted_nodes: Sequence[NodeIdent] | NodeIdent

要高亮的节点标识符。

highlight_css instance-attribute

highlight_css: str

用于高亮节点的 CSS。

title instance-attribute

title: str | None

图表的标题。

edge_labels instance-attribute

edge_labels: bool

是否在图表中包含边标签。

notes instance-attribute

notes: bool

是否在图表的节点上包含注释，默认为 true。

image_type instance-attribute

image_type: Literal['jpeg', 'png', 'webp', 'svg', 'pdf']

要生成的图像类型。如果未指定，则默认行为为 'jpeg'。

pdf_fit instance-attribute

pdf_fit: bool

当使用 image_type='pdf' 时，是否将图表适应 PDF 页面。

pdf_landscape instance-attribute

pdf_landscape: bool

当使用 image_type='pdf' 时，是否对 PDF 使用横向方向。

如果使用 pdf_fit，则此设置无效。

pdf_paper instance-attribute

pdf_paper: Literal[
    "letter",
    "legal",
    "tabloid",
    "ledger",
    "a0",
    "a1",
    "a2",
    "a3",
    "a4",
    "a5",
    "a6",
]

当使用 image_type='pdf' 时，PDF 的纸张尺寸。

background_color instance-attribute

background_color: str

图表的背景颜色。

如果为 None，则使用默认的透明背景。颜色值默认解释为十六进制颜色代码（并且不应带有前导 '#'），但您也可以通过在值前加上 '!' 来使用命名颜色。例如，有效的选择包括 background_color='!white' 或 background_color='FF0000'。

theme instance-attribute

theme: Literal['default', 'neutral', 'dark', 'forest']

图表的主题。默认为 'default'。

width instance-attribute

width: int

图表的宽度。

height instance-attribute

height: int

图表的高度。

scale instance-attribute

scale: Annotated[float, Ge(1), Le(3)]

图表的缩放比例。

缩放比例必须是 1 到 3 之间的数字，并且只有在设置了宽度和高度之一或两者时才能设置缩放比例。

httpx_client instance-attribute

httpx_client: Client

用于请求的 HTTPX 客户端，主要用于测试目的。

direction instance-attribute

direction: StateDiagramDirection

状态图的方向。

NodeIdent module-attribute

NodeIdent: TypeAlias = (
    "type[BaseNode[Any, Any, Any]] | BaseNode[Any, Any, Any] | str"
)

节点标识符的类型别名。

这可以是

• 节点实例（BaseNode 子类的实例）。
• 节点类（BaseNode 的子类）。
• 表示节点 ID 的字符串。