Skip to content

Commit

Permalink
Remove debug output; Generic runner base
Browse files Browse the repository at this point in the history
  • Loading branch information
nx10 committed Jul 19, 2024
1 parent 8c201d2 commit 3430e66
Showing 1 changed file with 11 additions and 9 deletions.
20 changes: 11 additions & 9 deletions src/styxgraph/__init__.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
""".. include:: ../../README.md""" # noqa: D415

import pathlib
from typing import Generic, TypeVar

from styxdefs import (
Execution,
Expand All @@ -10,6 +11,8 @@
Runner,
)

T = TypeVar("T", bound=Runner)


class _GraphExecution(Execution):
"""Graph execution."""
Expand Down Expand Up @@ -44,13 +47,13 @@ def output_file(self, local_file: str, optional: bool = False) -> OutputPathType


# Define a new runner
class GraphRunner(Runner):
class GraphRunner(Runner, Generic[T]):
"""Graph runner."""

def __init__(self, base: Runner) -> None:
def __init__(self, base: T) -> None:
"""Create a new GraphRunner."""
self.base = base
self.graph: list[tuple[str, list[InputPathType], list[OutputPathType]]] = []
self.base: T = base
self._graph: list[tuple[str, list[InputPathType], list[OutputPathType]]] = []

def start_execution(self, metadata: Metadata) -> Execution:
"""Start a new execution."""
Expand All @@ -63,23 +66,22 @@ def graph_append(
output_file: list[OutputPathType],
) -> None:
"""Append a node to the graph."""
print(f"Appending {metadata.name} with {input_file} and {output_file}")
self.graph.append((metadata.name, input_file, output_file))
self._graph.append((metadata.name, input_file, output_file))

def mermaid(self) -> str:
"""Generate a mermaid graph of the graph."""
connections: list[str] = []
inputs_lookup: dict[str, list[str]] = {}
outputs_lookup: dict[str, str] = {}
for id, inputs, outputs in self.graph:
for id, inputs, outputs in self._graph:
for input in inputs:
if input not in inputs_lookup:
inputs_lookup[input] = []
inputs_lookup[input].append(id)
root_output = outputs[0]
outputs_lookup[root_output] = id

for id, inputs, _ in self.graph:
for id, inputs, _ in self._graph:
for input in inputs:
for output, output_id in outputs_lookup.items():
if pathlib.Path(input).is_relative_to(
Expand All @@ -89,7 +91,7 @@ def mermaid(self) -> str:

# Generate mermaid
mermaid = "graph TD\n"
for id, _, _ in self.graph:
for id, _, _ in self._graph:
mermaid += f" {id}\n"
for connection in connections:
mermaid += f" {connection}\n"
Expand Down

0 comments on commit 3430e66

Please sign in to comment.