code2graph/graph_schema.py at main · VictorRaman/code2graph · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
"""Data structure skeletons for code graph nodes and edges."""

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from typing import Any


@dataclass
class Node:
    """Code graph node."""

    id: str
    type: str
    name: str
    path: str | None = None
    lineno: int | None = None
    end_lineno: int | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Convert the node to a serializable dictionary."""
        return asdict(self)


@dataclass
class Edge:
    """Code graph edge."""

    source: str
    target: str
    type: str
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Convert the edge to a serializable dictionary."""
        return asdict(self)