|
7 | 7 | from __future__ import annotations |
8 | 8 |
|
9 | 9 | import typing as t |
| 10 | +from contextlib import contextmanager |
10 | 11 | from itertools import cycle |
11 | 12 | from operator import itemgetter |
12 | 13 |
|
@@ -36,43 +37,35 @@ def _has_split_overlays(obj) -> TypeIs[HoloMap]: |
36 | 37 | return hasattr(obj, "_split_overlays") |
37 | 38 |
|
38 | 39 |
|
39 | | -class item_check: |
| 40 | +@contextmanager |
| 41 | +def item_check(enabled): |
40 | 42 | """Context manager to allow creating NdMapping types without |
41 | 43 | performing the usual item_checks, providing significant |
42 | 44 | speedups when there are a lot of items. Should only be |
43 | 45 | used when both keys and values are guaranteed to be the |
44 | 46 | right type, as is the case for many internal operations. |
45 | | -
|
46 | 47 | """ |
47 | | - |
48 | | - def __init__(self, enabled): |
49 | | - self.enabled = enabled |
50 | | - |
51 | | - def __enter__(self): |
52 | | - self._enabled = MultiDimensionalMapping._check_items |
53 | | - MultiDimensionalMapping._check_items = self.enabled |
54 | | - |
55 | | - def __exit__(self, exc_type, exc_val, exc_tb): |
56 | | - MultiDimensionalMapping._check_items = self._enabled |
| 48 | + prev = MultiDimensionalMapping._check_items |
| 49 | + try: |
| 50 | + MultiDimensionalMapping._check_items = enabled |
| 51 | + yield |
| 52 | + finally: |
| 53 | + MultiDimensionalMapping._check_items = prev |
57 | 54 |
|
58 | 55 |
|
59 | | -class sorted_context: |
| 56 | +@contextmanager |
| 57 | +def sorted_context(enabled): |
60 | 58 | """Context manager to temporarily disable sorting on NdMapping |
61 | 59 | types. Retains the current sort order, which can be useful as |
62 | 60 | an optimization on NdMapping instances where sort=True but the |
63 | 61 | items are already known to have been sorted. |
64 | | -
|
65 | 62 | """ |
66 | | - |
67 | | - def __init__(self, enabled): |
68 | | - self.enabled = enabled |
69 | | - |
70 | | - def __enter__(self): |
71 | | - self._enabled = MultiDimensionalMapping.sort |
72 | | - MultiDimensionalMapping.sort = self.enabled |
73 | | - |
74 | | - def __exit__(self, exc_type, exc_val, exc_tb): |
75 | | - MultiDimensionalMapping.sort = self._enabled # ty:ignore[invalid-assignment] |
| 63 | + prev = MultiDimensionalMapping.sort |
| 64 | + try: |
| 65 | + MultiDimensionalMapping.sort = enabled |
| 66 | + yield |
| 67 | + finally: |
| 68 | + MultiDimensionalMapping.sort = prev # ty:ignore[invalid-assignment] |
76 | 69 |
|
77 | 70 |
|
78 | 71 | class MultiDimensionalMapping(Dimensioned): |
|
0 commit comments