diff --git a/changelog.md b/changelog.md index 762168aa..8036a3d8 100644 --- a/changelog.md +++ b/changelog.md @@ -9,6 +9,7 @@ ## New features * Allow passing a `neo4j.Driver` instance as input to `from_neo4j`, in which case the driver will be used internally to fetch the graph data using a simple query +* The `rel_dfs` parameter of `from_dfs` is now optional, allowing for loading a graph with only nodes and no relationships ## Bug fixes diff --git a/python-wrapper/src/neo4j_viz/pandas.py b/python-wrapper/src/neo4j_viz/pandas.py index b07d9c39..3efaddd1 100644 --- a/python-wrapper/src/neo4j_viz/pandas.py +++ b/python-wrapper/src/neo4j_viz/pandas.py @@ -27,13 +27,19 @@ def _parse_validation_error(e: ValidationError, entity_type: type[BaseModel]) -> def _from_dfs( - node_dfs: Optional[DFS_TYPE], - rel_dfs: DFS_TYPE, + node_dfs: Optional[DFS_TYPE] = None, + rel_dfs: Optional[DFS_TYPE] = None, node_radius_min_max: Optional[tuple[float, float]] = (3, 60), rename_properties: Optional[dict[str, str]] = None, dropna: bool = False, ) -> VisualizationGraph: - relationships = _parse_relationships(rel_dfs, rename_properties=rename_properties, dropna=dropna) + if node_dfs is None and rel_dfs is None: + raise ValueError("At least one of `node_dfs` or `rel_dfs` must be provided") + + if rel_dfs is None: + relationships = [] + else: + relationships = _parse_relationships(rel_dfs, rename_properties=rename_properties, dropna=dropna) if node_dfs is None: has_size = False @@ -124,8 +130,8 @@ def _parse_relationships( def from_dfs( - node_dfs: Optional[DFS_TYPE], - rel_dfs: DFS_TYPE, + node_dfs: Optional[DFS_TYPE] = None, + rel_dfs: Optional[DFS_TYPE] = None, node_radius_min_max: Optional[tuple[float, float]] = (3, 60), ) -> VisualizationGraph: """ @@ -137,11 +143,12 @@ def from_dfs( Parameters ---------- - node_dfs: Optional[Union[DataFrame, Iterable[DataFrame]]] + node_dfs: Optional[Union[DataFrame, Iterable[DataFrame]]], optional DataFrame or iterable of DataFrames containing node data. If None, the nodes will be created from the source and target node ids in the rel_dfs. - rel_dfs: Union[DataFrame, Iterable[DataFrame]] + rel_dfs: Optional[Union[DataFrame, Iterable[DataFrame]]], optional DataFrame or iterable of DataFrames containing relationship data. + If None, no relationships will be created. node_radius_min_max : tuple[float, float], optional Minimum and maximum node radius. To avoid tiny or huge nodes in the visualization, the node sizes are scaled to fit in the given range. diff --git a/python-wrapper/tests/test_pandas.py b/python-wrapper/tests/test_pandas.py index 7da45bf1..b29f2dbf 100644 --- a/python-wrapper/tests/test_pandas.py +++ b/python-wrapper/tests/test_pandas.py @@ -190,3 +190,39 @@ def test_rel_errors() -> None: match=r"Error for relationship column 'caption_size' with provided input '-300.0'. Reason: Input should be greater than 0", ): from_dfs(nodes, relationships) + + +def test_from_dfs_no_rels() -> None: + nodes = [ + DataFrame( + { + "id": [0], + "caption": ["A"], + "size": [1337], + "color": "#FF0000", + } + ), + DataFrame( + { + "id": [1], + "caption": ["B"], + "size": [42], + "color": "#FF0000", + } + ), + ] + VG = from_dfs(nodes, [], node_radius_min_max=(42, 1337)) + + assert len(VG.nodes) == 2 + + assert VG.nodes[0].id == 0 + assert VG.nodes[0].caption == "A" + assert VG.nodes[0].size == 1337 + assert VG.nodes[0].color == Color("#ff0000") + + assert VG.nodes[1].id == 1 + assert VG.nodes[1].caption == "B" + assert VG.nodes[1].size == 42 + assert VG.nodes[0].color == Color("#ff0000") + + assert len(VG.relationships) == 0