diff --git a/examples/jaffle_shop/docs/README.md b/examples/jaffle_shop/docs/README.md index f5460b5..f4a2a4c 100644 --- a/examples/jaffle_shop/docs/README.md +++ b/examples/jaffle_shop/docs/README.md @@ -23,20 +23,27 @@ flowchart TB ```mermaid %%{init: {"flowchart": {"defaultRenderer": "elk"}} }%% flowchart TB + subgraph analytics - analytics.finance.kpis(finance.kpis) - analytics.kpis(kpis) + + subgraph finance + analytics.finance.kpis(kpis) + end + + analytics.kpis(kpis) end + subgraph core - core.customers(customers) - core.orders(orders) + core.customers(customers) + core.orders(orders) end + subgraph staging - staging.customers(customers) - staging.orders(orders) - staging.payments(payments) + staging.customers(customers) + staging.orders(orders) + staging.payments(payments) end core.orders --> analytics.finance.kpis diff --git a/lea/views/dag.py b/lea/views/dag.py index 93b7952..be3b199 100644 --- a/lea/views/dag.py +++ b/lea/views/dag.py @@ -3,10 +3,11 @@ import collections import graphlib import io -import itertools import lea +FOUR_SPACES = " " + class DAGOfViews(graphlib.TopologicalSorter, collections.UserDict): def __init__(self, views: list[lea.views.View]): @@ -47,6 +48,35 @@ def _list_descendants(node): return list(_list_descendants(node)) + def _build_nested_schema(self, deps): + nested_schema = {} + + for path in deps: + current_level = nested_schema + for part in path: + if part not in current_level: + current_level[part] = {} + current_level = current_level[part] + + return nested_schema + + def _to_mermaid_subgraphs(self, out, nested_schema: dict): + def output_subgraph(schema: str, values: dict, prefix: str = ""): + out.write(f"\n{FOUR_SPACES}subgraph {schema}\n") + for value in sorted(values.keys()): + sub_values = values[value] + path = f"{prefix}.{schema}" if prefix else schema + full_path = f"{path}.{value}" + if sub_values: + output_subgraph(value, sub_values, path) + else: + out.write(f"{FOUR_SPACES*2}{full_path}({value})\n") + out.write(f"{FOUR_SPACES}end\n\n") + + for schema in sorted(nested_schema.keys()): + values = nested_schema[schema] + output_subgraph(schema, values) + def _to_mermaid_views(self): out = io.StringIO() out.write('%%{init: {"flowchart": {"defaultRenderer": "elk"}} }%%\n') @@ -54,18 +84,15 @@ def _to_mermaid_views(self): nodes = set(node for deps in self.dependencies.values() for node in deps) | set( self.dependencies.keys() ) - schema_nodes = itertools.groupby(sorted(nodes), lambda node: node[0]) - for schema, nodes in schema_nodes: - out.write(f" subgraph {schema}\n") - for _, *node in sorted(nodes): - node = ".".join(node) - out.write(f" {schema}.{node}({node})\n") - out.write(" end\n\n") + + nested_schema = self._build_nested_schema(nodes) + self._to_mermaid_subgraphs(out, nested_schema) + for dst, srcs in sorted(self.dependencies.items()): dst = ".".join(dst) for src in sorted(srcs): src = ".".join(src) - out.write(f" {src} --> {dst}\n") + out.write(f"{FOUR_SPACES}{src} --> {dst}\n") return out.getvalue() def _to_mermaid_schemas(self): @@ -77,10 +104,10 @@ def _to_mermaid_schemas(self): schema_dependencies.keys() ) for node in sorted(nodes): - out.write(f" {node}({node})\n") + out.write(f"{FOUR_SPACES}{node}({node})\n") for dst, srcs in sorted(schema_dependencies.items()): for src in sorted(srcs): - out.write(f" {src} --> {dst}\n") + out.write(f"{FOUR_SPACES}{src} --> {dst}\n") return out.getvalue() def to_mermaid(self, schemas_only=False): @@ -107,20 +134,27 @@ def to_mermaid(self, schemas_only=False): >>> print(dag.to_mermaid()) %%{init: {"flowchart": {"defaultRenderer": "elk"}} }%% flowchart TB + subgraph analytics - analytics.finance.kpis(finance.kpis) - analytics.kpis(kpis) + + subgraph finance + analytics.finance.kpis(kpis) + end + + analytics.kpis(kpis) end + subgraph core - core.customers(customers) - core.orders(orders) + core.customers(customers) + core.orders(orders) end + subgraph staging - staging.customers(customers) - staging.orders(orders) - staging.payments(payments) + staging.customers(customers) + staging.orders(orders) + staging.payments(payments) end core.orders --> analytics.finance.kpis