From cd0b901ecaa87ef9b3a35e6572a2779e42792bd1 Mon Sep 17 00:00:00 2001 From: Anke Koke Date: Wed, 8 Nov 2023 12:16:44 +0100 Subject: [PATCH] fix: remove private module use --- niceml/mkdocs/mdjob.py | 8 +++----- niceml/mkdocs/mdop.py | 8 +++----- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/niceml/mkdocs/mdjob.py b/niceml/mkdocs/mdjob.py index e6236b54..5d62e427 100644 --- a/niceml/mkdocs/mdjob.py +++ b/niceml/mkdocs/mdjob.py @@ -1,11 +1,9 @@ """Module for generating mkdocs str for jobs""" from typing import List -from dagster._core.definitions import NodeDefinition - from niceml.mkdocs.mdgraph import get_graph_md from niceml.mkdocs.mdop import get_md_op -from dagster import JobDefinition +from dagster import JobDefinition, OpDefinition def get_job_md(job: JobDefinition, include_graph: bool = True) -> str: @@ -17,13 +15,13 @@ def get_job_md(job: JobDefinition, include_graph: bool = True) -> str: graph_md = get_graph_md(job) if len(graph_md) > 0: job_md += graph_md + "\n\n" - op_list: List[NodeDefinition] = get_ops_from_job(job) + op_list: List[OpDefinition] = get_ops_from_job(job) for cur_op in op_list: job_md += get_md_op(cur_op) return job_md -def get_ops_from_job(job: JobDefinition) -> List[NodeDefinition]: +def get_ops_from_job(job: JobDefinition) -> List[OpDefinition]: """Returns all ops from job""" return job.all_node_defs diff --git a/niceml/mkdocs/mdop.py b/niceml/mkdocs/mdop.py index 92602a98..23f71195 100644 --- a/niceml/mkdocs/mdop.py +++ b/niceml/mkdocs/mdop.py @@ -1,13 +1,11 @@ """Module for generating markdown strings for dagster ops""" from typing import Dict, List -from dagster._core.definitions import NodeDefinition - from niceml.mkdocs.mdtable import get_md_table -from dagster import Field +from dagster import Field, OpDefinition -def get_md_op(op_def: NodeDefinition) -> str: +def get_md_op(op_def: OpDefinition) -> str: """generates markdown strings for dagster ops""" col_widths: List[int] = [80, 120] op_fields = get_op_fields(op_def) @@ -24,7 +22,7 @@ def get_md_op(op_def: NodeDefinition) -> str: return cur_md -def get_op_fields(op_def: NodeDefinition) -> Dict[str, Field]: +def get_op_fields(op_def: OpDefinition) -> Dict[str, Field]: """returns fields from OpDefinition""" try: return op_def.config_schema.config_type.fields