_plotutils.py

import matplotlib.pyplot as plt

import _graph as g


def _is_edges(items):
    # an edge list looks like [((x1, y1), (x2, y2)), ...] — each item's
    # first element is itself a point (tuple/list), not a number. a vertex
    # list looks like [(x1, y1), (x2, y2), ...] — each item's first element
    # is a plain number.
    return isinstance(items[0][0], (tuple, list))


def _plot_shape(ax, points):
    # unzip points into per-axis coordinate sequences: (x, y) for 2D points,
    # (x, y, z) for 3D points
    coords = list(zip(*points))

    # close the shape: append each axis's first value to the end of its
    # sequence so the last segment connects back to the first vertex
    coords = [axis_values + (axis_values[0],) for axis_values in coords]

    # draws connected line segments through the points, "o" marks each vertex
    ax.plot(*coords, marker="o")


def _plot_edges(ax, pairs):
    # pairs is a list of (point_a, point_b) edges; each is drawn as its own
    # line segment, so unlike _plot_shape() there's no implicit closing of
    # the loop
    for point_a, point_b in pairs:
        coords = list(zip(point_a, point_b))
        ax.plot(*coords, marker="o")


def plot(items):
    is_edges = _is_edges(items)
    first_point = items[0][0] if is_edges else items[0]
    dim = len(first_point)

    # matplotlib fixes an axes' projection at creation time, so the point
    # dimensionality has to be known before the axes exists
    ax = plt.figure().add_subplot(projection="3d" if dim == 3 else None)

    if is_edges:
        _plot_edges(ax, items)
    else:
        _plot_shape(ax, items)

    g.style_axes(ax, dim)
    g.render()