Visualize a pixi lockfile using Rerun.io

visualize_lockfile.py

# Run this with:
#    pixi exec -s py-rattler -s rerun-sdk -s networkx  python visualize_lockfile.py

from __future__ import annotations
import argparse
import random
import rerun as rr
import rerun.blueprint as rrb
from rerun.blueprint.archetypes.force_collision_radius import ForceCollisionRadius
from rerun.blueprint.archetypes.force_link import ForceLink
from rerun.blueprint.archetypes.force_many_body import ForceManyBody
from rerun.blueprint.archetypes.force_center import ForceCenter
from rerun.components.color import Color
from rattler import LockFile, Platform, CondaLockedBinaryPackage, PackageRecord
import networkx as nx

# Seed for reproducibility
random.seed(42)

def log_lock_file_graph(lockfile_path: str) -> None:
    # Load the lock file
    lock_file = LockFile.from_path(lockfile_path)

    # Get packages for the current platform
    packages = list(lock_file.default_environment().packages(platform=Platform.current()))
    package_records = [package.package_record for package in packages]

    # Create graph
    graph = PackageRecord.to_graph(package_records)

    # Find depth for every node
    roots = [node for node in graph.nodes if graph.in_degree(node) == 0]
    depths = {node: 0 for node in roots}

    for root in roots:
        for node, depth in nx.single_source_shortest_path_length(graph, root).items():
            depths[node] = max(depths.get(node, 0), depth)

    # Generate a color gradient based on depth
    max_depth = max(depths.values(), default=1)
    colors = []
    for node in graph.nodes:
        depth = depths.get(node, 0)
        # Map depth to a color (example: gradient from red to green)
        r = int((1 - depth / max_depth) * 255)
        g = int((depth / max_depth) * 255)
        colors.append(Color([r, g, 0]))

    # Filter for CondaLockedBinaryPackage instances
    conda_binary_packages = [
        package for package in packages
        if isinstance(package, CondaLockedBinaryPackage)
    ]

    # Generate nodes and edges
    nodes = []
    labels = []
    edges = []

    for package in conda_binary_packages:
        # Add the node (package name and version)
        nodes.append(package.name)
        labels.append(f"{package.name} {package.version}")

        # Add edges based on dependencies
        for dependency in package.package_record.depends:
            if dependency.startswith("__"):
                continue
            dependency_name = dependency.split()[0]  # Extract just the name from dependency spec
            edges.append((str(package.name), str(dependency_name)))

    rr.log("lock_file", rr.GraphEdges(edges, graph_type="directed"))
    rr.log("lock_file", rr.GraphNodes(nodes, labels=labels, colors=colors))


def log_dashboard() -> None:
    rr.send_blueprint(
        rrb.Blueprint(
            rrb.Grid(
                rrb.GraphView(
                    origin="lock_file",
                    name="Lock File Dependencies",
                    force_link=ForceLink(distance=100),
                    force_many_body=ForceManyBody(enabled=True, strength=-1500),
                    force_center=ForceCenter(enabled=True, strength=1),
                    force_collision_radius=ForceCollisionRadius(strength=2, enabled=True)
                ),
            )
        )
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Visualize the pixi lockfile as a graph")
    parser.add_argument(
        "lockfile",
        type=str,
        nargs="?",
        default="./pixi.lock",
        help="Path to the lock file (default: ./pixi.lock)"
    )
    rr.script_add_args(parser)
    args = parser.parse_args()

    rr.script_setup(args, "Pixi LockFile Graph")

    log_lock_file_graph(args.lockfile)
    log_dashboard()

    rr.script_teardown(args)


if __name__ == "__main__":
    main()

As the script says you can run this with:

pixi exec -s py-rattler -s rerun-sdk -s networkx  python visualize_lockfile.py