feat: add d2 export functionality (#33)

arguebuf/dump/__init__.py

@@ -8,12 +8,14 @@
 from ._dump_path import dump_file as file
 from ._dump_path import dump_file as folder
 from ._dump_protobuf import dump_protobuf as protobuf
+from ._dump_d2 import dump_d2 as d2
 
 __all__ = (
     "file",
     "folder",
     "networkx",
     "graphviz",
+    "d2",
     "aif",
     "protobuf",
     "dict",

arguebuf/dump/_dump_d2.py

@@ -0,0 +1,79 @@
+import textwrap
+import typing as t
+
+from arguebuf.model import Graph
+from arguebuf.model.edge import Edge
+from arguebuf.model.node import AtomNode, SchemeNode, AbstractNode
+from arguebuf.schemas.d2 import D2Graph, D2Edge, D2Node, D2Style
+
+
+def dump_d2(
+    graph: Graph,
+    atom_label: t.Optional[t.Callable[[AtomNode], str]] = None,
+    scheme_label: t.Optional[t.Callable[[SchemeNode], str]] = None,
+    max_nodes: t.Optional[int] = None,
+) -> t.Optional[D2Graph]:
+    if len(graph.nodes) > (max_nodes or 1000):
+        return None
+
+    d2_graph = D2Graph(
+        nodes=[],
+        edges=[],
+    )
+
+    for node in graph._atom_nodes.values():
+        _dump_node(
+            node,
+            d2_graph,
+            major_claim=graph.major_claim == node,
+            label_func=atom_label,
+        )
+
+    for node in graph._scheme_nodes.values():
+        _dump_node(
+            node,
+            d2_graph,
+            label_func=scheme_label,
+            major_claim=False,
+        )
+
+    for edge in graph._edges.values():
+        _dump_edge(edge, d2_graph)
+
+    return d2_graph
+
+
+def _dump_node(
+    node: AbstractNode,
+    g: D2Graph,
+    major_claim: bool,
+    label_func: t.Optional[t.Callable[[AbstractNode], str]] = None,
+) -> None:
+    label: str = label_func(node) if label_func else node.label
+    label = label.replace("\n", "")
+    if type(node) == AtomNode:
+        color = node.color(major_claim)
+    else:
+        color = node.color(major_claim=False)
+
+    nodeStyle: D2Style = D2Style(
+        font_color=color.fg,
+        stroke_width=2,
+        bold=False,
+        stroke=color.border,
+        fill=color.bg,
+    )
+
+    g.nodes.append(
+        D2Node(
+            id=node.id,
+            label=label,
+            shape="rectangle",
+            style=nodeStyle,
+        )
+    )
+
+
+def _dump_edge(edge: Edge, g: D2Graph) -> None:
+    """Submethod used to export Graph object g into D2 Graph format."""
+    g.edges.append(D2Edge(from_id=edge.source.id, to_id=edge.target.id))

arguebuf/load/_load_xaif.py

@@ -17,9 +17,7 @@
 
 
 def load_xaif(
-        obj: xaif.Graph,
-        name: t.Optional[str] = None,
-        config: Config = DefaultConfig
+    obj: xaif.Graph, name: t.Optional[str] = None, config: Config = DefaultConfig
 ) -> Graph:
     """
     Generate Graph structure from xAif argument graph file
@@ -44,10 +42,7 @@ def load_xaif(
     return g
 
 
-def scheme_from_xaif(
-        obj: xaif.AifNode,
-        config: Config
-) -> t.Optional[SchemeNode]:
+def scheme_from_xaif(obj: xaif.AifNode, config: Config) -> t.Optional[SchemeNode]:
     """Generate SchemeNode object from xAif Node object."""
 
     aif_type = obj["type"]
@@ -73,10 +68,7 @@ def scheme_from_xaif(
     return None
 
 
-def atom_from_xaif(
-        obj: xaif.AifNode,
-        config: Config
-) -> AtomNode:
+def atom_from_xaif(obj: xaif.AifNode, config: Config) -> AtomNode:
     """Generate AtomNode object from xAif Node object."""
     timestamp = pendulum.now()
 
@@ -88,9 +80,7 @@ def atom_from_xaif(
 
 
 def edge_from_xaif(
-        obj: xaif.AifEdge,
-        nodes: t.Mapping[str, AbstractNode],
-        config: Config
+    obj: xaif.AifEdge, nodes: t.Mapping[str, AbstractNode], config: Config
 ) -> t.Optional[Edge]:
     """Generate Edge object from xAif Edge format."""
     source_id = obj.get("fromID")

arguebuf/schemas/d2.py

@@ -0,0 +1,63 @@
+import typing as t
+
+
+class D2Style:
+    def __init__(
+        self, font_color: str, bold: bool, stroke: str, stroke_width: int, fill: str
+    ):
+        self.font_color: str = font_color
+        self.bold = bold
+        self.stroke = stroke
+        self.stroke_width = stroke_width
+        self.fill = fill
+
+
+class D2Node:
+    def __init__(self, id: str, label: str, shape: str, style: D2Style):
+        self.id = id
+        self.label = label
+        self.shape = shape
+        self.style = style
+
+
+class D2Edge:
+    def __init__(self, from_id: str, to_id: str):
+        self.from_id = from_id
+        self.to_id = to_id
+
+
+class D2Graph:
+    def __init__(self, nodes: t.List[D2Node], edges: t.List[D2Edge]):
+        self.nodes = nodes
+        self.edges = edges
+
+    def __str__(self):
+        str = ""
+        for node in self.nodes:
+            str += self.__write_node(node)
+        for edge in self.edges:
+            str += self.__write_edge(edge)
+        return str
+
+    def __write_node(self, node: D2Node) -> str:
+        str = ""
+        label = node.label.replace('"', '\\"')
+        str += node.id + ": {\n"
+        str += "label: " + label + "\n"
+        str += "shape: " + node.shape + "\n"
+        str += "style: " + self.__write_style(node.style) + "\n}\n"
+        return str
+
+    def __write_style(self, style: D2Style) -> str:
+        str = "{\n"
+        str += "font-color: " + '"' + style.font_color + '"' + "\n"
+        str += "bold: " + style.bold.__str__() + "\n"
+        str += "stroke: " + '"' + style.stroke + '"' + "\n"
+        str += "stroke-width: " + style.stroke_width.__str__() + "\n"
+        str += "fill: " + '"' + style.fill + '"' + "\n}"
+        return str
+
+    def __write_edge(self, edge: D2Edge) -> str:
+        str = ""
+        str += edge.from_id + " -> " + edge.to_id + "\n"
+        return str