neo4j-viz 0.2.6__py3-none-any.whl → 0.3.0__py3-none-any.whl

neo4j_viz/colors.py

@@ -1,12 +1,34 @@
 from collections.abc import Iterable
+from enum import Enum
 from typing import Any, Union
 
+import enum_tools
 from pydantic_extra_types.color import ColorType
 
 ColorsType = Union[dict[Any, ColorType], Iterable[ColorType]]
 
+
+@enum_tools.documentation.document_enum
+class ColorSpace(Enum):
+    """
+    Describes the type of color space used by a color palette.
+    """
+
+    DISCRETE = "discrete"
+    """
+    This category describes a color palette that is a collection of different colors that are not necessarily related to
+    each other. Discrete color spaces are suitable for categorical data, where each unique category is represented by a
+    different color.
+    """
+    CONTINUOUS = "continuous"
+    """
+    This category describes a color palette that is a range/gradient of colors between two or more colors. Continuous
+    color spaces are suitable for continuous data (typically floats), where values can change smoothly.
+    """
+
+
 # Comes from https://neo4j.design/40a8cff71/p/5639c0-color/t/page-5639c0-79109681-33
-neo4j_colors = [
+NEO4J_COLORS_DISCRETE = [
     "#FFDF81",
     "#C990C0",
     "#F79767",
@@ -20,3 +42,262 @@ neo4j_colors = [
     "#DA7294",
     "#579380",
 ]
+_NEO4J_COLORS_CONTINUOUS_BASE = ["#FFDF81", "#C990C0"]
+NEO4J_COLORS_CONTINUOUS = [
+    (255, 223, 129),
+    (255, 223, 129),
+    (255, 222, 129),
+    (254, 222, 130),
+    (254, 222, 130),
+    (254, 221, 130),
+    (254, 221, 130),
+    (254, 221, 131),
+    (253, 221, 131),
+    (253, 220, 131),
+    (253, 220, 131),
+    (253, 220, 132),
+    (252, 219, 132),
+    (252, 219, 132),
+    (252, 219, 132),
+    (252, 218, 133),
+    (252, 218, 133),
+    (251, 218, 133),
+    (251, 217, 133),
+    (251, 217, 134),
+    (251, 217, 134),
+    (251, 216, 134),
+    (250, 216, 134),
+    (250, 216, 135),
+    (250, 216, 135),
+    (250, 215, 135),
+    (249, 215, 135),
+    (249, 215, 136),
+    (249, 214, 136),
+    (249, 214, 136),
+    (249, 214, 136),
+    (248, 213, 137),
+    (248, 213, 137),
+    (248, 213, 137),
+    (248, 212, 137),
+    (248, 212, 138),
+    (247, 212, 138),
+    (247, 212, 138),
+    (247, 211, 138),
+    (247, 211, 139),
+    (247, 211, 139),
+    (246, 210, 139),
+    (246, 210, 139),
+    (246, 210, 140),
+    (246, 209, 140),
+    (245, 209, 140),
+    (245, 209, 140),
+    (245, 208, 141),
+    (245, 208, 141),
+    (245, 208, 141),
+    (244, 208, 141),
+    (244, 207, 142),
+    (244, 207, 142),
+    (244, 207, 142),
+    (244, 206, 142),
+    (243, 206, 143),
+    (243, 206, 143),
+    (243, 205, 143),
+    (243, 205, 143),
+    (243, 205, 144),
+    (242, 204, 144),
+    (242, 204, 144),
+    (242, 204, 144),
+    (242, 203, 145),
+    (241, 203, 145),
+    (241, 203, 145),
+    (241, 203, 145),
+    (241, 202, 146),
+    (241, 202, 146),
+    (240, 202, 146),
+    (240, 201, 146),
+    (240, 201, 147),
+    (240, 201, 147),
+    (240, 200, 147),
+    (239, 200, 147),
+    (239, 200, 148),
+    (239, 199, 148),
+    (239, 199, 148),
+    (238, 199, 148),
+    (238, 199, 149),
+    (238, 198, 149),
+    (238, 198, 149),
+    (238, 198, 149),
+    (237, 197, 150),
+    (237, 197, 150),
+    (237, 197, 150),
+    (237, 196, 150),
+    (237, 196, 150),
+    (236, 196, 151),
+    (236, 195, 151),
+    (236, 195, 151),
+    (236, 195, 151),
+    (236, 194, 152),
+    (235, 194, 152),
+    (235, 194, 152),
+    (235, 194, 152),
+    (235, 193, 153),
+    (234, 193, 153),
+    (234, 193, 153),
+    (234, 192, 153),
+    (234, 192, 154),
+    (234, 192, 154),
+    (233, 191, 154),
+    (233, 191, 154),
+    (233, 191, 155),
+    (233, 190, 155),
+    (233, 190, 155),
+    (232, 190, 155),
+    (232, 190, 156),
+    (232, 189, 156),
+    (232, 189, 156),
+    (231, 189, 156),
+    (231, 188, 157),
+    (231, 188, 157),
+    (231, 188, 157),
+    (231, 187, 157),
+    (230, 187, 158),
+    (230, 187, 158),
+    (230, 186, 158),
+    (230, 186, 158),
+    (230, 186, 159),
+    (229, 186, 159),
+    (229, 185, 159),
+    (229, 185, 159),
+    (229, 185, 160),
+    (229, 184, 160),
+    (228, 184, 160),
+    (228, 184, 160),
+    (228, 183, 161),
+    (228, 183, 161),
+    (227, 183, 161),
+    (227, 182, 161),
+    (227, 182, 162),
+    (227, 182, 162),
+    (227, 181, 162),
+    (226, 181, 162),
+    (226, 181, 163),
+    (226, 181, 163),
+    (226, 180, 163),
+    (226, 180, 163),
+    (225, 180, 164),
+    (225, 179, 164),
+    (225, 179, 164),
+    (225, 179, 164),
+    (225, 178, 165),
+    (224, 178, 165),
+    (224, 178, 165),
+    (224, 177, 165),
+    (224, 177, 166),
+    (223, 177, 166),
+    (223, 177, 166),
+    (223, 176, 166),
+    (223, 176, 167),
+    (223, 176, 167),
+    (222, 175, 167),
+    (222, 175, 167),
+    (222, 175, 168),
+    (222, 174, 168),
+    (222, 174, 168),
+    (221, 174, 168),
+    (221, 173, 169),
+    (221, 173, 169),
+    (221, 173, 169),
+    (220, 173, 169),
+    (220, 172, 170),
+    (220, 172, 170),
+    (220, 172, 170),
+    (220, 171, 170),
+    (219, 171, 171),
+    (219, 171, 171),
+    (219, 170, 171),
+    (219, 170, 171),
+    (219, 170, 171),
+    (218, 169, 172),
+    (218, 169, 172),
+    (218, 169, 172),
+    (218, 168, 172),
+    (218, 168, 173),
+    (217, 168, 173),
+    (217, 168, 173),
+    (217, 167, 173),
+    (217, 167, 174),
+    (216, 167, 174),
+    (216, 166, 174),
+    (216, 166, 174),
+    (216, 166, 175),
+    (216, 165, 175),
+    (215, 165, 175),
+    (215, 165, 175),
+    (215, 164, 176),
+    (215, 164, 176),
+    (215, 164, 176),
+    (214, 164, 176),
+    (214, 163, 177),
+    (214, 163, 177),
+    (214, 163, 177),
+    (213, 162, 177),
+    (213, 162, 178),
+    (213, 162, 178),
+    (213, 161, 178),
+    (213, 161, 178),
+    (212, 161, 179),
+    (212, 160, 179),
+    (212, 160, 179),
+    (212, 160, 179),
+    (212, 159, 180),
+    (211, 159, 180),
+    (211, 159, 180),
+    (211, 159, 180),
+    (211, 158, 181),
+    (211, 158, 181),
+    (210, 158, 181),
+    (210, 157, 181),
+    (210, 157, 182),
+    (210, 157, 182),
+    (209, 156, 182),
+    (209, 156, 182),
+    (209, 156, 183),
+    (209, 155, 183),
+    (209, 155, 183),
+    (208, 155, 183),
+    (208, 155, 184),
+    (208, 154, 184),
+    (208, 154, 184),
+    (208, 154, 184),
+    (207, 153, 185),
+    (207, 153, 185),
+    (207, 153, 185),
+    (207, 152, 185),
+    (207, 152, 186),
+    (206, 152, 186),
+    (206, 151, 186),
+    (206, 151, 186),
+    (206, 151, 187),
+    (205, 151, 187),
+    (205, 150, 187),
+    (205, 150, 187),
+    (205, 150, 188),
+    (205, 149, 188),
+    (204, 149, 188),
+    (204, 149, 188),
+    (204, 148, 189),
+    (204, 148, 189),
+    (204, 148, 189),
+    (203, 147, 189),
+    (203, 147, 190),
+    (203, 147, 190),
+    (203, 146, 190),
+    (202, 146, 190),
+    (202, 146, 191),
+    (202, 146, 191),
+    (202, 145, 191),
+    (202, 145, 191),
+    (201, 145, 192),
+    (201, 144, 192),
+    (201, 144, 192),
+]

neo4j_viz/gds.py

@@ -6,7 +6,7 @@ from typing import Optional
 import pandas as pd
 from graphdatascience import Graph, GraphDataScience
 
-from .pandas import from_dfs
+from .pandas import _from_dfs
 from .visualization_graph import VisualizationGraph
 
 
@@ -35,6 +35,11 @@ def from_gds(
     """
     Create a VisualizationGraph from a GraphDataScience object and a Graph object.
 
+    All `additional_node_properties` will be included in the visualization graph.
+    If the properties are named as the fields of the `Node` class, they will be included as top level fields of the
+    created `Node` objects. Otherwise, they will be included in the `properties` dictionary.
+    Additionally, a new "labels" node property will be added, containing the node labels of the node.
+
     Parameters
     ----------
     gds : GraphDataScience
@@ -75,9 +80,13 @@ def from_gds(
 
     node_props_df = pd.concat(node_dfs.values(), ignore_index=True, axis=0).drop_duplicates()
     if size_property is not None:
+        if "size" in actual_node_properties and size_property != "size":
+            node_props_df.rename(columns={"size": "__size"}, inplace=True)
         node_props_df.rename(columns={size_property: "size"}, inplace=True)
 
     for lbl, df in node_dfs.items():
+        if "labels" in actual_node_properties:
+            df.rename(columns={"labels": "__labels"}, inplace=True)
         df["labels"] = lbl
 
     node_lbls_df = pd.concat([df[["id", "labels"]] for df in node_dfs.values()], ignore_index=True, axis=0)
@@ -88,4 +97,4 @@ def from_gds(
     rel_df = _rel_df(gds, G)
     rel_df.rename(columns={"sourceNodeId": "source", "targetNodeId": "target"}, inplace=True)
 
-    return from_dfs(node_df, rel_df, node_radius_min_max=node_radius_min_max)
+    return _from_dfs(node_df, rel_df, node_radius_min_max=node_radius_min_max, rename_properties={"__size": "size"})

neo4j_viz/gql_create.py

@@ -0,0 +1,356 @@
+import re
+import uuid
+from typing import Any, Optional
+
+from neo4j_viz import Node, Relationship, VisualizationGraph
+
+
+def _parse_value(value_str: str) -> Any:
+    value_str = value_str.strip()
+    if not value_str:
+        return None
+
+    # Parse map
+    if value_str.startswith("{") and value_str.endswith("}"):
+        inner = value_str[1:-1].strip()
+        result = {}
+        depth = 0
+        in_string = None
+        start_idx = 0
+        for i, ch in enumerate(inner):
+            if in_string is None:
+                if ch in ["'", '"']:
+                    in_string = ch
+                elif ch in ["{", "["]:
+                    depth += 1
+                elif ch in ["}", "]"]:
+                    depth -= 1
+                elif ch == "," and depth == 0:
+                    segment = inner[start_idx:i].strip()
+                    if ":" not in segment:
+                        return None
+                    k, v = segment.split(":", 1)
+                    k = k.strip().strip("'\"")
+                    result[k] = _parse_value(v)
+                    start_idx = i + 1
+            else:
+                if ch == in_string:
+                    in_string = None
+
+        if inner[start_idx:]:
+            segment = inner[start_idx:].strip()
+            if ":" not in segment:
+                return None
+            k, v = segment.split(":", 1)
+            k = k.strip().strip("'\"")
+            result[k] = _parse_value(v)
+
+        return result
+
+    # Parse list
+    if value_str.startswith("[") and value_str.endswith("]"):
+        inner = value_str[1:-1].strip()
+        items = []
+        depth = 0
+        in_string = None
+        start_idx = 0
+        for i, ch in enumerate(inner):
+            if in_string is None:
+                if ch in ["'", '"']:
+                    in_string = ch
+                elif ch in ["{", "["]:
+                    depth += 1
+                elif ch in ["}", "]"]:
+                    depth -= 1
+                elif ch == "," and depth == 0:
+                    items.append(_parse_value(inner[start_idx:i]))
+                    start_idx = i + 1
+            else:
+                if ch == in_string:
+                    in_string = None
+
+        if inner[start_idx:]:
+            items.append(_parse_value(inner[start_idx:]))
+
+        return items
+
+    # Parse boolean, float, int, or string
+    if re.match(r"^-?\d+$", value_str):
+        return int(value_str)
+    if re.match(r"^-?\d+\.\d+$", value_str):
+        return float(value_str)
+    if value_str.lower() == "true":
+        return True
+    if value_str.lower() == "false":
+        return False
+    if value_str.lower() == "null":
+        return None
+
+    return value_str.strip("'\"")
+
+
+def _parse_prop_str(
+    query: str, prop_str: str, prop_start: int, top_level_keys: set[str]
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    top_level: dict[str, Any] = {}
+    props: dict[str, Any] = {}
+    depth = 0
+    in_string = None
+    start_idx = 0
+    for i, ch in enumerate(prop_str):
+        if in_string is None:
+            if ch in ["'", '"']:
+                in_string = ch
+            elif ch in ["{", "["]:
+                depth += 1
+            elif ch in ["}", "]"]:
+                depth -= 1
+            elif ch == "," and depth == 0:
+                pair = prop_str[start_idx:i].strip()
+                if ":" not in pair:
+                    snippet = _get_snippet(query, prop_start + start_idx)
+                    raise ValueError(f"Property syntax error near: `{snippet}`.")
+                k, v = pair.split(":", 1)
+                k = k.strip().strip("'\"")
+
+                if k in top_level_keys:
+                    top_level[k] = _parse_value(v)
+                else:
+                    props[k] = _parse_value(v)
+
+                start_idx = i + 1
+        else:
+            if ch == in_string:
+                in_string = None
+
+    if prop_str[start_idx:]:
+        pair = prop_str[start_idx:].strip()
+        if ":" not in pair:
+            snippet = _get_snippet(query, prop_start + start_idx)
+            raise ValueError(f"Property syntax error near: `{snippet}`.")
+        k, v = pair.split(":", 1)
+        k = k.strip().strip("'\"")
+
+        if k in top_level_keys:
+            top_level[k] = _parse_value(v)
+        else:
+            props[k] = _parse_value(v)
+
+    return top_level, props
+
+
+def _parse_labels_and_props(
+    query: str, s: str, top_level_keys: set[str]
+) -> tuple[Optional[str], dict[str, Any], dict[str, Any]]:
+    prop_match = re.search(r"\{(.*)\}", s)
+    prop_str = ""
+    if prop_match:
+        prop_str = prop_match.group(1)
+        prop_start = query.index(prop_str, query.index(s))
+        s = s[: prop_match.start()].strip()
+    alias_labels = re.split(r"[:&]", s)
+    raw_alias = alias_labels[0].strip()
+    final_alias = raw_alias if raw_alias else None
+
+    if prop_str:
+        top_level, props = _parse_prop_str(query, prop_str, prop_start, top_level_keys)
+    else:
+        top_level = {}
+        props = {}
+
+    label_list = [lbl.strip() for lbl in alias_labels[1:]]
+    if "labels" in props:
+        props["__labels"] = props["labels"]
+    props["labels"] = sorted(label_list)
+
+    return final_alias, top_level, props
+
+
+def _get_snippet(q: str, idx: int, context: int = 15) -> str:
+    start = max(0, idx - context)
+    end = min(len(q), idx + context)
+
+    return q[start:end].replace("\n", " ")
+
+
+def from_gql_create(
+    query: str,
+    size_property: Optional[str] = None,
+    node_caption: Optional[str] = "labels",
+    relationship_caption: Optional[str] = "type",
+    node_radius_min_max: Optional[tuple[float, float]] = (3, 60),
+) -> VisualizationGraph:
+    """
+    Parse a GQL CREATE query and return a VisualizationGraph object representing the graph it creates.
+
+    All node and relationship properties will be included in the visualization graph.
+    If the properties are named as the fields of the `Node` or `Relationship` classes, they will be included as
+    top level fields of the respective objects. Otherwise, they will be included in the `properties` dictionary.
+    Additionally, a "labels" property will be added for nodes and a "type" property for relationships.
+
+    Please note that this function is not a full GQL parser, it only handles CREATE queries that do not contain
+    other clauses like MATCH, WHERE, RETURN, etc, or any Cypher function calls.
+    It also does not handle all possible GQL syntax, but it should work for most common cases.
+    For more complex cases, we recommend using a Neo4j database and the `from_neo4j` method.
+
+    Parameters
+    ----------
+    query : str
+        The GQL CREATE query to parse
+    size_property : str, optional
+        Property to use for node size, by default None.
+    node_caption : str, optional
+        Property to use as the node caption, by default the node labels will be used.
+    relationship_caption : str, optional
+        Property to use as the relationship caption, by default the relationship type will be used.
+    node_radius_min_max : tuple[float, float], optional
+        Minimum and maximum node radius, by default (3, 60).
+        To avoid tiny or huge nodes in the visualization, the node sizes are scaled to fit in the given range.
+    """
+
+    query = query.strip()
+    if not re.match(r"(?i)^create\b", query):
+        raise ValueError("Query must begin with 'CREATE' (case insensitive).")
+
+    query = re.sub(r"(?i)^create\s*", "", query, count=1).rstrip(";").strip()
+    parts = []
+    paren_level = 0
+    bracket_level = 0
+    current: list[str] = []
+    for i, char in enumerate(query):
+        if char == "(":
+            paren_level += 1
+        elif char == ")":
+            paren_level -= 1
+            if paren_level < 0:
+                snippet = _get_snippet(query, i)
+                raise ValueError(f"Unbalanced parentheses near: `{snippet}`.")
+        if char == "[":
+            bracket_level += 1
+        elif char == "]":
+            bracket_level -= 1
+            if bracket_level < 0:
+                snippet = _get_snippet(query, i)
+                raise ValueError(f"Unbalanced square brackets near: `{snippet}`.")
+        if char == "," and paren_level == 0 and bracket_level == 0:
+            parts.append("".join(current).strip())
+            current = []
+        else:
+            current.append(char)
+
+    parts.append("".join(current).strip())
+    if paren_level != 0:
+        snippet = _get_snippet(query, len(query) - 1)
+        raise ValueError(f"Unbalanced parentheses near: `{snippet}`.")
+    if bracket_level != 0:
+        snippet = _get_snippet(query, len(query) - 1)
+        raise ValueError(f"Unbalanced square brackets near: `{snippet}`.")
+
+    node_pattern = re.compile(r"^\(([^)]*)\)$")
+    rel_pattern = re.compile(r"^\(([^)]*)\)-\s*\[\s*:(\w+)\s*(\{[^}]*\})?\s*\]->\(([^)]*)\)$")
+
+    node_top_level_keys = set(Node.model_fields.keys())
+    node_top_level_keys.remove("id")
+
+    rel_top_level_keys = set(Relationship.model_fields.keys())
+    rel_top_level_keys.remove("id")
+    rel_top_level_keys.remove("source")
+    rel_top_level_keys.remove("target")
+
+    nodes = []
+    relationships = []
+    alias_to_id = {}
+    anonymous_count = 0
+
+    for part in parts:
+        node_m = node_pattern.match(part)
+        if node_m:
+            alias_labels_props = node_m.group(1).strip()
+            alias, top_level, props = _parse_labels_and_props(query, alias_labels_props, node_top_level_keys)
+            if not alias:
+                alias = f"_anon_{anonymous_count}"
+                anonymous_count += 1
+            if alias not in alias_to_id:
+                alias_to_id[alias] = str(uuid.uuid4())
+            nodes.append(Node(id=alias_to_id[alias], **top_level, properties=props))
+
+            continue
+
+        rel_m = rel_pattern.match(part)
+        if rel_m:
+            left_node = rel_m.group(1).strip()
+            right_node = rel_m.group(4).strip()
+
+            # Parse left node pattern
+            left_alias, left_top_level, left_props = _parse_labels_and_props(query, left_node, node_top_level_keys)
+            if not left_alias:
+                left_alias = f"_anon_{anonymous_count}"
+                anonymous_count += 1
+                if left_alias not in alias_to_id:
+                    alias_to_id[left_alias] = str(uuid.uuid4())
+                nodes.append(Node(id=alias_to_id[left_alias], **left_top_level, properties=left_props))
+            elif left_alias not in alias_to_id:
+                snippet = _get_snippet(query, query.index(left_node))
+                raise ValueError(f"Relationship references unknown node alias: '{left_alias}' near: `{snippet}`.")
+
+            # Parse right node pattern
+            right_alias, right_top_level, right_props = _parse_labels_and_props(query, right_node, node_top_level_keys)
+            if not right_alias:
+                right_alias = f"_anon_{anonymous_count}"
+                anonymous_count += 1
+                if right_alias not in alias_to_id:
+                    alias_to_id[right_alias] = str(uuid.uuid4())
+                nodes.append(Node(id=alias_to_id[right_alias], **right_top_level, properties=right_props))
+            elif right_alias not in alias_to_id:
+                snippet = _get_snippet(query, query.index(right_node))
+                raise ValueError(f"Relationship references unknown node alias: '{right_alias}' near: `{snippet}`.")
+
+            rel_id = str(uuid.uuid4())
+            rel_type = rel_m.group(2).replace(":", "").strip()
+            rel_props_str = rel_m.group(3) or ""
+            if rel_props_str:
+                inner_str = rel_props_str.strip("{}").strip()
+                prop_start = query.index(inner_str, query.index(inner_str))
+                top_level, props = _parse_prop_str(query, inner_str, prop_start, rel_top_level_keys)
+            else:
+                top_level = {}
+                props = {}
+            if "type" in props:
+                props["__type"] = props["type"]
+            props["type"] = rel_type
+            relationships.append(
+                Relationship(
+                    id=rel_id,
+                    source=alias_to_id[left_alias],
+                    target=alias_to_id[right_alias],
+                    **top_level,
+                    properties=props,
+                )
+            )
+            continue
+
+        snippet = part[:30]
+        raise ValueError(f"Invalid element in CREATE near: `{snippet}`.")
+
+    if size_property is not None:
+        for node in nodes:
+            node.size = node.properties.get(size_property)
+    if node_caption is not None:
+        for node in nodes:
+            if node_caption == "labels":
+                if len(node.properties["labels"]) > 0:
+                    node.caption = ":".join([label for label in node.properties["labels"]])
+            else:
+                node.caption = str(node.properties.get(node_caption))
+    if relationship_caption is not None:
+        for rel in relationships:
+            if relationship_caption == "type":
+                rel.caption = rel.properties["type"]
+            else:
+                rel.caption = str(rel.properties.get(relationship_caption))
+
+    VG = VisualizationGraph(nodes=nodes, relationships=relationships)
+    if (node_radius_min_max is not None) and (size_property is not None):
+        VG.resize_nodes(node_radius_min_max=node_radius_min_max)
+
+    return VG