Generate visualizations of the graph

cag/cli.py

@@ -1,13 +1,20 @@
+import importlib
+import inspect
 import os
 import platform
 import zipfile
 from pathlib import Path
 from subprocess import PIPE, run
 
 import typer
-from slugify import slugify
+from pyvis.network import Network
 from rich import print
+from rich.console import Console
+from slugify import slugify
+
+from cag.framework import GraphCreatorBase
 
+console = Console()
 app = typer.Typer()
 
 
@@ -61,5 +68,88 @@ def start_project():
     print(f"[green] Creating scaffold in {project_dir}[/green]")
 
 
+@app.command()
+def visualize(python_file: Path) -> Path:
+    """
+    Visualize the graph of the collections and relations defined in a Python file.
+
+    Parameters:
+    -----------
+    python_file : Path
+        The path of the Python file containing a GraphCreatorBase class that define the collections
+        and relations to be visualized.
+
+    Raises:
+    -------
+    typer.Abort:
+        If the specified `python_file` does not exist or is not a file.
+
+   Returns:
+    --------
+    Path
+        The path of the HTML file containing the generated graph.
+        The function saves the generated graph to an HTML file and prints the path of the file.
+    """
+    if not python_file.exists():
+        print(f"{python_file} not found!")
+        raise typer.Abort()
+    if not python_file.is_file():
+        print(f"{python_file} is not a file")
+        raise typer.Abort()
+    spec = importlib.util.spec_from_file_location(
+        python_file.stem, python_file
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    net = Network(height="750px", width="100%")
+    net.barnes_hut()
+
+    subclasses = []
+    for name, obj in inspect.getmembers(module, inspect.isclass):
+        if (
+            issubclass(obj, module.GraphCreatorBase)
+            and obj != module.GraphCreatorBase
+        ):
+            subclasses.append(obj)
+
+    if len(subclasses) == 0:
+        print("No suitable classes found!")
+        raise typer.Abort()
+    with console.status("[bold green]Working...") as status:
+        for subclass in subclasses:  # type: GraphCreatorBase
+            for ed_def in subclass._edge_definitions:
+                for collection in set(
+                    ed_def["from_collections"] + ed_def["to_collections"]
+                ):
+                    collection_name = GraphCreatorBase.get_collection_name(
+                        collection
+                    )
+                    net.add_node(collection_name, label=collection_name)
+                    console.log(f"Process {collection}")
+
+                from_collections = [
+                    GraphCreatorBase.get_collection_name(m)
+                    for m in ed_def["from_collections"]
+                ]
+                to_collections = [
+                    GraphCreatorBase.get_collection_name(m)
+                    for m in ed_def["to_collections"]
+                ]
+                relation_name = GraphCreatorBase.get_collection_name(
+                    ed_def["relation"]
+                )
+                net.add_edge(
+                    *from_collections, *to_collections, label=relation_name
+                )
+
+        diagram_file = python_file.parent.joinpath(
+            python_file.stem + "_diagram.html"
+        )
+        status.update("Save file")
+        net.save_graph(str(diagram_file))
+    print(f"File saved to {diagram_file}")
+    return diagram_file
+
+
 if __name__ == "__main__":
     app()

pyproject.toml

@@ -34,7 +34,10 @@ dependencies = ['dataclasses>=0.6',
                     'python-arango>=7.4.1',
                     'pyArango>=2.0.1',
                     'tomli>=2.0.1',
-                    'transformers>=4.26.1'
+                    'transformers>=4.26.1',
+                    'python-slugify',
+                    'rich>=13.3.2',
+                    'beautifulsoup4>=4.11.2'
                       ]
 requires-python = ">=3.8"
 

requirements.txt

@@ -14,4 +14,5 @@ tomli>=2.0.1
 typer~=0.4.2
 python-slugify~=6.1.2
 rich~=12.6.0
-transformers>=4.26.1
+transformers>=4.26.1
+python-slugify

tests/test_24-generate-visualizations-of-the-graph.py

@@ -0,0 +1,82 @@
+from bs4 import BeautifulSoup
+from pathlib import Path
+
+from cag.cli import visualize
+from cag.framework import GraphCreatorBase
+from cag.graph_elements.nodes import GenericOOSNode, Field
+from cag.graph_elements.relations import GenericEdge
+
+
+class CollectionA(GenericOOSNode):
+    _name = "CollectionA"
+    _fields = {"value": Field(), "value2": Field(), **GenericOOSNode._fields}
+
+
+class CollectionB(GenericOOSNode):
+    _name = "CollectionB"
+    _fields = {"value": Field(), **GenericOOSNode._fields}
+
+
+class CollectionC(GenericOOSNode):
+    _name = "CollectionC"
+    _fields = {"value": Field(), **GenericOOSNode._fields}
+
+
+class HasRelation(GenericEdge):
+    _fields = GenericEdge._fields
+
+
+class HasAnotherRelation(GenericEdge):
+    _fields = GenericEdge._fields
+
+
+class SampleGraphCreator(GraphCreatorBase):
+    _name = "SampleGraphCreator"
+    _description = "Sample Graph"
+
+    _edge_definitions = [
+        {
+            "relation": "HasRelation",
+            "from_collections": [CollectionA],
+            "to_collections": [CollectionB],
+        },
+        {
+            "relation": "HasAnotherRelation",
+            "from_collections": [CollectionC],
+            "to_collections": [CollectionC],
+        },
+    ]
+
+    def init_graph(self):
+        pass
+
+
+class Test24:
+    def test_generate(self):
+        """
+        Test the cli `visualize` function by generating a diagram HTML file and checking if it exists and contains
+        valid HTML.
+
+        This test calls the `visualize` function with the path of this Python file, which contain a Graphcreator
+        derived from the `GraphCreatorBase` class. The `visualize` function generates an HTML file that shows the
+        relationships between the collections and relations defined in those class.
+
+        This test first checks if the generated HTML file exists. If the file exists, it opens the file and parses it
+        using the `BeautifulSoup` class from the `bs4` module. If the parsing succeeds, the test passes. Otherwise,
+        the test fails.
+
+        Note: The purpose of this test is to check if the file contains anything useful at all, not to verify its exact
+        structure.
+
+        After the test, the generated HTML file is deleted to clean up the test environment.
+        """
+        html_file = visualize(Path(__file__).absolute())
+        assert html_file.exists()
+        # This test may fail in the future
+        # the purpose is to check if the file contains anything useful at all
+        with open(html_file) as f:
+            BeautifulSoup(f, "html.parser")
+            # if the call doesn't throw any exceptions, the test passes
+            assert True
+        # clean up
+        html_file.unlink()

tests/test_graph_creator/test_improve_edge_definitions.py

@@ -50,14 +50,14 @@ class SampleGraphCreator(GraphCreatorBase):
             "to_collections": [CollectionB],
         },
         {
-            "relation": "HasAnotherRelation",
+            "relation": HasAnotherRelation,
             "from_collections": [CollectionC],
             "to_collections": [CollectionC],
         },
         {  # old style
-            "relation": "HasAnotherAnotherRelation",
-            "from_collections": ["CollectionA"],
-            "to_collections": ["CollectionA"],
+            "relation": HasAnotherAnotherRelation,
+            "from_collections": [CollectionA],
+            "to_collections": [CollectionA],
         },
     ]
 

tests/test_graph_creator/test_simple_init.py

@@ -34,7 +34,7 @@ class HasAnotherRelation(GenericEdge):
 
 class SampleGraphCreator(GraphCreatorBase):
     _name = "SampleGraphCreator"
-    _description = "Graph based on the DLR elib corpus"
+    _description = "Sample Graph"
 
     _edge_definitions = [
         {