127 lines
4.2 KiB
Python
127 lines
4.2 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import statistics
|
||
|
|
from fractions import Fraction
|
||
|
|
from typing import Callable, Generic, Iterable, Sequence, TypeVar
|
||
|
|
|
||
|
|
from rich.color import Color
|
||
|
|
from rich.console import Console, ConsoleOptions, RenderResult
|
||
|
|
from rich.measure import Measurement
|
||
|
|
from rich.segment import Segment
|
||
|
|
from rich.style import Style
|
||
|
|
|
||
|
|
from textual.renderables._blend_colors import blend_colors
|
||
|
|
|
||
|
|
T = TypeVar("T", int, float)
|
||
|
|
|
||
|
|
SummaryFunction = Callable[[Sequence[T]], float]
|
||
|
|
|
||
|
|
|
||
|
|
class Sparkline(Generic[T]):
|
||
|
|
"""A sparkline representing a series of data.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
data: The sequence of data to render.
|
||
|
|
width: The width of the sparkline/the number of buckets to partition the data into.
|
||
|
|
min_color: The color of values equal to the min value in data.
|
||
|
|
max_color: The color of values equal to the max value in data.
|
||
|
|
summary_function: Function that will be applied to each bucket.
|
||
|
|
"""
|
||
|
|
|
||
|
|
BARS = "▁▂▃▄▅▆▇█"
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
data: Sequence[T],
|
||
|
|
*,
|
||
|
|
width: int | None,
|
||
|
|
min_color: Color = Color.from_rgb(0, 255, 0),
|
||
|
|
max_color: Color = Color.from_rgb(255, 0, 0),
|
||
|
|
summary_function: SummaryFunction[T] = max,
|
||
|
|
) -> None:
|
||
|
|
self.data: Sequence[T] = data
|
||
|
|
self.width = width
|
||
|
|
self.min_color = Style.from_color(min_color)
|
||
|
|
self.max_color = Style.from_color(max_color)
|
||
|
|
self.summary_function: SummaryFunction[T] = summary_function
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _buckets(cls, data: list[T], num_buckets: int) -> Iterable[Sequence[T]]:
|
||
|
|
"""Partition ``data`` into ``num_buckets`` buckets. For example, the data
|
||
|
|
[1, 2, 3, 4] partitioned into 2 buckets is [[1, 2], [3, 4]].
|
||
|
|
|
||
|
|
Args:
|
||
|
|
data: The data to partition.
|
||
|
|
num_buckets: The number of buckets to partition the data into.
|
||
|
|
"""
|
||
|
|
bucket_step = Fraction(len(data), num_buckets)
|
||
|
|
for bucket_no in range(num_buckets):
|
||
|
|
start = int(bucket_step * bucket_no)
|
||
|
|
end = int(bucket_step * (bucket_no + 1))
|
||
|
|
partition = data[start:end]
|
||
|
|
if partition:
|
||
|
|
yield partition
|
||
|
|
|
||
|
|
def __rich_console__(
|
||
|
|
self, console: Console, options: ConsoleOptions
|
||
|
|
) -> RenderResult:
|
||
|
|
width = self.width or options.max_width
|
||
|
|
len_data = len(self.data)
|
||
|
|
if len_data == 0:
|
||
|
|
yield Segment("▁" * width, self.min_color)
|
||
|
|
return
|
||
|
|
if len_data == 1:
|
||
|
|
yield Segment("█" * width, self.max_color)
|
||
|
|
return
|
||
|
|
|
||
|
|
minimum, maximum = min(self.data), max(self.data)
|
||
|
|
extent = maximum - minimum or 1
|
||
|
|
|
||
|
|
buckets = tuple(self._buckets(list(self.data), num_buckets=width))
|
||
|
|
|
||
|
|
bucket_index = 0.0
|
||
|
|
bars_rendered = 0
|
||
|
|
step = len(buckets) / width
|
||
|
|
summary_function = self.summary_function
|
||
|
|
min_color, max_color = self.min_color.color, self.max_color.color
|
||
|
|
assert min_color is not None
|
||
|
|
assert max_color is not None
|
||
|
|
while bars_rendered < width:
|
||
|
|
partition = buckets[int(bucket_index)]
|
||
|
|
partition_summary = summary_function(partition)
|
||
|
|
height_ratio = (partition_summary - minimum) / extent
|
||
|
|
bar_index = int(height_ratio * (len(self.BARS) - 1))
|
||
|
|
bar_color = blend_colors(min_color, max_color, height_ratio)
|
||
|
|
bars_rendered += 1
|
||
|
|
bucket_index += step
|
||
|
|
yield Segment(self.BARS[bar_index], Style.from_color(bar_color))
|
||
|
|
|
||
|
|
def __rich_measure__(
|
||
|
|
self, console: "Console", options: "ConsoleOptions"
|
||
|
|
) -> Measurement:
|
||
|
|
return Measurement(self.width or options.max_width, 1)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
console = Console()
|
||
|
|
|
||
|
|
def last(l: Sequence[T]) -> T:
|
||
|
|
return l[-1]
|
||
|
|
|
||
|
|
funcs: Sequence[SummaryFunction[int]] = (
|
||
|
|
min,
|
||
|
|
max,
|
||
|
|
last,
|
||
|
|
statistics.median,
|
||
|
|
statistics.mean,
|
||
|
|
)
|
||
|
|
nums = [10, 2, 30, 60, 45, 20, 7, 8, 9, 10, 50, 13, 10, 6, 5, 4, 3, 7, 20]
|
||
|
|
console.print(f"data = {nums}\n")
|
||
|
|
for f in funcs:
|
||
|
|
console.print(
|
||
|
|
f"{f.__name__}:\t",
|
||
|
|
Sparkline(nums, width=12, summary_function=f),
|
||
|
|
end="",
|
||
|
|
)
|
||
|
|
console.print("\n")
|