From ed961db221eb8ca3253c2bab5a0bd9804c6636e2 Mon Sep 17 00:00:00 2001 From: Anke Koke Date: Wed, 11 Oct 2023 16:07:44 +0200 Subject: [PATCH] fix: generate mkdocs graphs with multiple dependencies --- niceml/mkdocs/mdgraph.py | 10 ++++++++-- niceml/mkdocs/mdjob.py | 2 +- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/niceml/mkdocs/mdgraph.py b/niceml/mkdocs/mdgraph.py index 872b20ee..459f0714 100644 --- a/niceml/mkdocs/mdgraph.py +++ b/niceml/mkdocs/mdgraph.py @@ -1,5 +1,5 @@ """Module for generating a graph in mkdocs""" -from dagster import JobDefinition +from dagster import JobDefinition, DependencyDefinition, MultiDependencyDefinition def get_graph_md(job: JobDefinition) -> str: @@ -8,7 +8,13 @@ def get_graph_md(job: JobDefinition) -> str: graph_str = "" for key, value in deps.items(): for _, val2 in value.items(): - graph_str += f" {val2.node} --> {key.name};\n" + if isinstance(val2, DependencyDefinition): + graph_str += f" {val2.node} --> {key.name};\n" + elif isinstance(val2, MultiDependencyDefinition): + for dependency in val2.dependencies: + graph_str += f" {dependency.node} --> {key.name};\n" + else: + raise AttributeError("'val2' is not of expected type.") if len(graph_str) == 0: return "" diff --git a/niceml/mkdocs/mdjob.py b/niceml/mkdocs/mdjob.py index 858e4537..e6236b54 100644 --- a/niceml/mkdocs/mdjob.py +++ b/niceml/mkdocs/mdjob.py @@ -1,7 +1,7 @@ """Module for generating mkdocs str for jobs""" from typing import List -from dagster.core.definitions import NodeDefinition +from dagster._core.definitions import NodeDefinition from niceml.mkdocs.mdgraph import get_graph_md from niceml.mkdocs.mdop import get_md_op