|
| 1 | +""" |
| 2 | +Benchmark comprehensions. |
| 3 | +
|
| 4 | +Author: Carl Meyer |
| 5 | +""" |
| 6 | + |
| 7 | +from dataclasses import dataclass |
| 8 | +from enum import Enum |
| 9 | +from typing import Iterable, List, Optional |
| 10 | + |
| 11 | +import pyperf |
| 12 | + |
| 13 | + |
| 14 | +class WidgetKind(Enum): |
| 15 | + BIG = 1 |
| 16 | + SMALL = 2 |
| 17 | + |
| 18 | + |
| 19 | +@dataclass |
| 20 | +class Widget: |
| 21 | + widget_id: int |
| 22 | + creator_id: int |
| 23 | + derived_widget_ids: List[int] |
| 24 | + kind: WidgetKind |
| 25 | + has_knob: bool |
| 26 | + has_spinner: bool |
| 27 | + |
| 28 | + |
| 29 | +class WidgetTray: |
| 30 | + def __init__(self, owner_id: int, widgets: List[Widget]) -> None: |
| 31 | + self.owner_id = owner_id |
| 32 | + self.sorted_widgets: List[Widget] = [] |
| 33 | + self._add_widgets(widgets) |
| 34 | + |
| 35 | + def _any_knobby(self, widgets: Iterable[Optional[Widget]]) -> bool: |
| 36 | + return any(w.has_knob for w in widgets if w) |
| 37 | + |
| 38 | + def _is_big_spinny(self, widget: Widget) -> bool: |
| 39 | + return widget.kind == WidgetKind.BIG and widget.has_spinner |
| 40 | + |
| 41 | + def _add_widgets(self, widgets: List[Widget]) -> None: |
| 42 | + # sort order: mine first, then any widgets with derived knobby widgets in order of |
| 43 | + # number derived, then other widgets in order of number derived, and we exclude |
| 44 | + # big spinny widgets entirely |
| 45 | + widgets = [w for w in widgets if not self._is_big_spinny(w)] |
| 46 | + id_to_widget = {w.widget_id: w for w in widgets} |
| 47 | + id_to_derived = { |
| 48 | + w.widget_id: [id_to_widget.get(dwid) for dwid in w.derived_widget_ids] |
| 49 | + for w in widgets |
| 50 | + } |
| 51 | + sortable_widgets = [ |
| 52 | + ( |
| 53 | + w.creator_id == self.owner_id, |
| 54 | + self._any_knobby(id_to_derived[w.widget_id]), |
| 55 | + len(id_to_derived[w.widget_id]), |
| 56 | + w.widget_id, |
| 57 | + ) |
| 58 | + for w in widgets |
| 59 | + ] |
| 60 | + sortable_widgets.sort() |
| 61 | + self.sorted_widgets = [id_to_widget[sw[-1]] for sw in sortable_widgets] |
| 62 | + |
| 63 | + |
| 64 | +def make_some_widgets() -> List[Widget]: |
| 65 | + widget_id = 0 |
| 66 | + widgets = [] |
| 67 | + for creator_id in range(3): |
| 68 | + for kind in WidgetKind: |
| 69 | + for has_knob in [True, False]: |
| 70 | + for has_spinner in [True, False]: |
| 71 | + derived = [w.widget_id for w in widgets[::creator_id + 1]] |
| 72 | + widgets.append( |
| 73 | + Widget( |
| 74 | + widget_id, creator_id, derived, kind, has_knob, has_spinner |
| 75 | + ) |
| 76 | + ) |
| 77 | + widget_id += 1 |
| 78 | + assert len(widgets) == 24 |
| 79 | + return widgets |
| 80 | + |
| 81 | + |
| 82 | +def bench_comprehensions(loops: int) -> float: |
| 83 | + range_it = range(loops) |
| 84 | + widgets = make_some_widgets() |
| 85 | + t0 = pyperf.perf_counter() |
| 86 | + for _ in range_it: |
| 87 | + tray = WidgetTray(1, widgets) |
| 88 | + assert len(tray.sorted_widgets) == 18 |
| 89 | + return pyperf.perf_counter() - t0 |
| 90 | + |
| 91 | + |
| 92 | +if __name__ == "__main__": |
| 93 | + runner = pyperf.Runner() |
| 94 | + runner.metadata["description"] = "Benchmark comprehensions" |
| 95 | + runner.bench_time_func("comprehensions", bench_comprehensions) |
0 commit comments