Skip to content

Commit

Permalink
Eliminate nr-io-graphviz dependency and update databind to min `4…
Browse files Browse the repository at this point in the history
….5.0` (#267)

* Eliminate `nr-io-graphviz` dependency and update `databind` to min `4.5.0`

* fix unused ignore

* use networkx type stubs

* fix DiGraph not subscriptable at runtime

* hack for networkx stubs
  • Loading branch information
NiklasRosenstein authored Aug 2, 2024
1 parent ad00bd6 commit 69f2895
Show file tree
Hide file tree
Showing 5 changed files with 178 additions and 14 deletions.
6 changes: 6 additions & 0 deletions .changelog/_unreleased.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,12 @@ type = "refactor"
description = "move `kraken.std.git.gitignore.parser` to `kraken.common.gitignore`"
author = "@NiklasRosenstein"

[[entries]]
id = "bbdd557d-b4c2-4ba6-9adc-66c18f51a44c"
type = "hygiene"
description = "Eliminate `nr-io-graphviz` dependency and update `databind` to min `4.5.0`"
author = "@NiklasRosenstein"

[[entries]]
id = "0a74ab07-f895-48e8-8842-6c5382ba8402"
type = "refactor"
Expand Down
9 changes: 2 additions & 7 deletions kraken-build/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -37,14 +37,13 @@ kraken = "kraken.core.cli.main:main"
# ------------

[tool.poetry.dependencies]
databind-json = "^4.2.5"
databind = "^4.5.0"
dataclasses = {version = "^0.6", python = "<3.7"}
deprecated = "^1.2.13"
dill = ">=0.3.8,<0.4.0"
httpx = "^0.27.0"
keyring = "^25.0.0"
networkx = "^3.1"
nr-io-graphviz = "^0.1.1"
packaging = "^23.1"
pex = "^2.1.156"
python = ">=3.10"
Expand All @@ -64,6 +63,7 @@ types-Deprecated = "^1.2.9"
types-requests = "^2.28.0"
pytest-xdist = {version = "^3.5.0", extras = ["psutil"]}
mitmproxy = "^10.2.4"
types-networkx = "^3.2.1.20240703"

# Slap configuration
# ------------------
Expand All @@ -89,11 +89,6 @@ warn_unreachable = true
warn_unused_ignores = true
enable_error_code = "ignore-without-code, possibly-undefined"

[[tool.mypy.overrides]]
ignore_errors = true
ignore_missing_imports = true
module = "networkx.*"

[tool.ruff]
line-length = 120

Expand Down
158 changes: 158 additions & 0 deletions kraken-build/src/kraken/common/graphviz.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,158 @@
from __future__ import annotations

import http.server
import logging
import subprocess as sp
import tempfile
import webbrowser
from pathlib import Path
from typing import overload

import typing as t

logger = logging.getLogger(__name__)


class GraphvizWriter:
"""A helper class to write Dotviz files."""

def __init__(self, out: t.TextIO, indent: str = "\t") -> None:
"""
:param out: The output file to write to.
:param indent: The string to use for every level of indendation.
"""

self._out = out
self._indent = indent
self._level = 0
self._edge_type: list[str] = []

@property
def _line_prefix(self) -> str:
return self._level * self._indent

def _escape(self, name: str) -> str:
# TODO (@NiklasRosenstein): Improve escaping logic.
if "\n" in name:
raise ValueError("Cannot have newline (contained in {name!r})")
chars = " ,.:#$/&"
if any(c in name for c in chars):
name = f'"{name}"'
return name

def _write_attrs(self, attrs: dict[str, str | None]) -> None:
safe_attrs = {key: value for key, value in attrs.items() if value is not None}
if safe_attrs:
self._out.write("[")
self._out.write(" ".join(f"{self._escape(key)}={self._escape(value)}" for key, value in safe_attrs.items()))
self._out.write("]")

def _write_scope(self, keyword: str, name: str | None, attrs: dict[str, str | None]) -> None:
self._out.write(f"{self._line_prefix}{keyword} ")
if name is not None:
self._out.write(self._escape(name) + " ")
self._out.write("{\n")
self._level += 1
for key, value in attrs.items():
if value is not None:
self._out.write(f"{self._line_prefix}{self._escape(key)}={self._escape(value)};\n")

def graph(self, name: str | None = None, **attrs: str | None) -> None:
"""Open a `graph{}` block. Close it by calling :meth:`end`."""
self._write_scope("graph", name, attrs)
self._edge_type.append("--")

def digraph(self, name: str | None = None, **attrs: str | None) -> None:
"""Open a `digraph{}` block. Close it by calling :meth:`end`."""
self._write_scope("digraph", name, attrs)
self._edge_type.append("->")

def subgraph(self, name: str | None = None, **attrs: str | None) -> None:
"""Open a `subgraph{}` block. Close it by calling :meth:`end`."""
self._write_scope("subgraph", name, attrs)
self._edge_type.append(self._edge_type[-1])

def end(self) -> None:
"""Close a previouly opened block. Raises an :class:`AssertionError` if called too many times."""
assert self._level >= 1, "called end() too many times"
self._level -= 1
self._edge_type.pop()
self._out.write(self._line_prefix + "}\n")

def set_node_style(self, **attrs: str | None) -> None:
"""Sets the node style in the current block."""
self._out.write(self._line_prefix + "node ")
self._write_attrs(attrs)
self._out.write(";\n")

def node(self, node_id: str, **attrs: str | None) -> None:
"""Draw a node in the current context."""
self._out.write(self._line_prefix + self._escape(node_id))
if attrs:
self._write_attrs(attrs)
self._out.write(";\n")

def edge(self, source: str | t.Sequence[str], target: str | t.Sequence[str], **attrs: str | None) -> None:
"""Draw one or multiple edges in the current contect from source to target. Specifying multiple
nodes on either side will generate the cross product of edges between all nodes."""
if isinstance(source, str):
source = [source]
if isinstance(target, str):
target = [target]

if not source or not target:
raise ValueError("edge needs at least one source and at least one target")

def _write_nodes(nodes: t.Sequence[str]) -> None:
if len(nodes) == 1:
self._out.write(self._escape(nodes[0]))
else:
self._out.write("{")
self._out.write(" ".join(self._escape(node_id) for node_id in nodes))
self._out.write("}")

# TODO (@NiklasRosenstein): Does GraphViz support a syntax for multiple nodes on the left?
for node_id in source:
self._out.write(self._line_prefix)
_write_nodes([node_id])
self._out.write(f" {self._edge_type[-1]} ")
_write_nodes(target)
self._write_attrs(attrs)
self._out.write(";\n")


@overload
def render(graphviz_code: str, format: str, algorithm: str = ...) -> bytes:
"""Renders the *graphviz_code* to an image file of the specified *format*. The default format is `"dot"`."""


@overload
def render(graphviz_code: str, format: str, algorithm: str = ..., *, output_file: Path) -> None:
"""Renders the *graphviz_code* to a file."""


def render(graphviz_code: str, format: str, algorithm: str = "dot", *, output_file: Path | None = None) -> None | bytes:
command = [algorithm, f"-T{format}"]
if output_file is not None:
command += ["-o", str(output_file)]
try:
process = sp.run(command, input=graphviz_code.encode(), check=True, capture_output=True)
except sp.CalledProcessError as exc:
logger.error("%s: %s", exc, exc.stderr.decode())
raise
return process.stdout


def render_to_browser(graphviz_code: str, algorithm: str = "dot") -> None:
"""Renders the *graphviz_code* to an SVG file and opens it in the webbrowser. Blocks until the
browser opened the page."""

with tempfile.TemporaryDirectory() as tempdir:
svg_file = Path(tempdir) / "graph.svg"
render(graphviz_code, "svg", algorithm, output_file=svg_file)
server = http.server.HTTPServer(
("", 0),
lambda *args: http.server.SimpleHTTPRequestHandler(*args, directory=tempdir),
)
webbrowser.open(f"http://localhost:{server.server_port}/graph.svg")
server.handle_request()
4 changes: 2 additions & 2 deletions kraken-build/src/kraken/core/cli/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
from pathlib import Path
from typing import Any, NoReturn

from nr.io.graphviz.render import render_to_browser
from nr.io.graphviz.writer import GraphvizWriter
from kraken.common.graphviz import render_to_browser
from kraken.common.graphviz import GraphvizWriter

from kraken.common import (
BuildscriptMetadata,
Expand Down
15 changes: 10 additions & 5 deletions kraken-build/src/kraken/core/system/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,8 @@ def __init__(self, context: Context, populate: bool = True, parent: TaskGraph |
self._context = context

# Nodes have the form {'data': _Node} and edges have the form {'data': _Edge}.
self._digraph = DiGraph()
# NOTE: DiGraph is not runtime-subscriptable.
self._digraph: DiGraph[Address] = DiGraph()

# Keep track of task execution results.
self._results: dict[Address, TaskStatus] = {}
Expand Down Expand Up @@ -198,7 +199,7 @@ def _remove_nodes_keep_transitive_edges(self, nodes: Iterable[Address]) -> None:
)
self._digraph.remove_node(addr)

def _get_ready_graph(self) -> DiGraph:
def _get_ready_graph(self) -> DiGraph[Address]:
"""Updates the ready graph. Remove all ok tasks (successful or skipped) and any non-strict dependencies
(edges) on failed tasks."""

Expand All @@ -225,7 +226,7 @@ def set_non_strict_edge_for_removal(u: Address, v: Address) -> None:
else:
set_non_strict_edge_for_removal(failed_task_path, out_task_path)

return restricted_view(self._digraph, self._ok_tasks, removable_edges)
return cast("DiGraph[Address]", restricted_view(self._digraph, self._ok_tasks, removable_edges)) # type: ignore[no-untyped-call]

# Public API

Expand Down Expand Up @@ -380,7 +381,8 @@ def tasks(

tasks = (self.get_task(addr) for addr in self._digraph)
if goals:
tasks = (t for t in tasks if self._digraph.out_degree(t.address) == 0)
# HACK: Cast because of https://github.com/python/typeshed/pull/12472
tasks = (t for t in tasks if cast(int, self._digraph.out_degree(t.address)) == 0)
if pending:
tasks = (t for t in tasks if t.address not in self._results)
if failed:
Expand Down Expand Up @@ -494,7 +496,10 @@ def ready(self) -> list[Task]:

ready_graph = self._get_ready_graph()
root_set = (
node for node in ready_graph.nodes if ready_graph.in_degree(node) == 0 and node not in self._results
# HACK: Cast because of https://github.com/python/typeshed/pull/12472
node
for node in ready_graph.nodes
if cast(int, ready_graph.in_degree(node)) == 0 and node not in self._results
)
tasks = [self.get_task(addr) for addr in root_set]
if not tasks:
Expand Down

0 comments on commit 69f2895

Please sign in to comment.