diff --git a/cag/cli.py b/cag/cli.py index 396f7ac..eeda6ff 100644 --- a/cag/cli.py +++ b/cag/cli.py @@ -1,3 +1,5 @@ +import importlib +import inspect import os import platform import zipfile @@ -5,9 +7,14 @@ 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() diff --git a/pyproject.toml b/pyproject.toml index 4aa836d..b4f5043 100644 --- a/pyproject.toml +++ b/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" diff --git a/requirements.txt b/requirements.txt index fefd9f2..5f03306 100644 --- a/requirements.txt +++ b/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 \ No newline at end of file +transformers>=4.26.1 +python-slugify diff --git a/tests/test_24-generate-visualizations-of-the-graph.py b/tests/test_24-generate-visualizations-of-the-graph.py new file mode 100644 index 0000000..542bd93 --- /dev/null +++ b/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() diff --git a/tests/test_graph_creator/test_improve_edge_definitions.py b/tests/test_graph_creator/test_improve_edge_definitions.py index eb91bc2..07614cb 100644 --- a/tests/test_graph_creator/test_improve_edge_definitions.py +++ b/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], }, ] diff --git a/tests/test_graph_creator/test_simple_init.py b/tests/test_graph_creator/test_simple_init.py index 4c60d51..436762d 100644 --- a/tests/test_graph_creator/test_simple_init.py +++ b/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 = [ {