chromatic-python 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,1677 @@
1
+ __all__ = [
2
+ 'CSI',
3
+ 'Color',
4
+ 'ColorStr',
5
+ 'SgrParameter',
6
+ 'SgrSequence',
7
+ 'ansicolor24Bit',
8
+ 'ansicolor4Bit',
9
+ 'ansicolor8Bit',
10
+ 'colorbytes',
11
+ 'get_ansi_type',
12
+ 'hsl_gradient',
13
+ 'randcolor',
14
+ 'rgb2ansi_color_esc',
15
+ 'rgb_luma_transform',
16
+ 'SGR_RESET',
17
+ ]
18
+
19
+ import math
20
+ import operator as op
21
+ import os
22
+ import random
23
+ from collections import Counter
24
+ from collections.abc import Buffer
25
+ from copy import deepcopy
26
+ from ctypes import byref, c_ulong, windll
27
+ from enum import IntEnum
28
+ from functools import lru_cache
29
+ from types import MappingProxyType
30
+ from typing import (
31
+ Callable, cast, Final, Generator, Iterable, Iterator, Literal, Mapping, Optional, Self,
32
+ Sequence, SupportsIndex, SupportsInt, TypedDict, TypeVar, Union
33
+ )
34
+
35
+ import numpy as np
36
+
37
+ from .colorconv import *
38
+ from .._typing import (
39
+ AnsiColorAlias, ColorDictKeys, Float3Tuple, Int3Tuple, is_matching_typed_dict, RGBVectorLike
40
+ )
41
+
42
+ os.system('')
43
+
44
+ CSI: Final[bytes] = b'['
45
+ SGR_RESET: Final[str] = ''
46
+
47
+ # ansi color global lookups
48
+ # ansi 4bit {color code (int) ==> (key, RGB)}
49
+ _ANSI16C_I2KV = cast(
50
+ dict[int, tuple[ColorDictKeys, Int3Tuple]],
51
+ {v: (k, ansi_4bit_to_rgb(v))
52
+ for x in (
53
+ zip(
54
+ ('fg', 'bg'),
55
+ (j, j + 10))
56
+ for i in (30, 90)
57
+ for j in range(i, i + 8))
58
+ for (k, v) in x})
59
+
60
+ # ansi 4bit {(key, RGB) ==> color code (int)}
61
+ _ANSI16C_KV2I = {v: k for k, v in _ANSI16C_I2KV.items()}
62
+
63
+ # ansi 4bit standard color range
64
+ _ANSI16C_STD = frozenset(x for i in (30, 40) for x in range(i, i + 8))
65
+
66
+ # ansi 4bit bright color range
67
+ _ANSI16C_BRIGHT = frozenset(_ANSI16C_I2KV.keys() - _ANSI16C_STD)
68
+
69
+ # ansi 8bit {color code (bytes) ==> color dict key (str)}
70
+ _ANSI256_B2KEY = {b'38': 'fg', b'48': 'bg'}
71
+
72
+ # ansi 8bit {color dict key (str) ==> color code (int)}
73
+ _ANSI256_KEY2I = {v: int(k) for k, v in _ANSI256_B2KEY.items()}
74
+
75
+
76
+ # see also: https://en.wikipedia.org/wiki/ANSI_escape_code#SGR
77
+ # int enum {sgr parameter name ==> sgr code (int)}
78
+ class SgrParameter(IntEnum):
79
+ RESET = 0
80
+ BOLD = 1
81
+ FAINT = 2
82
+ ITALICS = 3
83
+ SINGLE_UNDERLINE = 4
84
+ SLOW_BLINK = 5
85
+ RAPID_BLINK = 6
86
+ NEGATIVE = 7
87
+ CONCEALED_CHARS = 8
88
+ CROSSED_OUT = 9
89
+ PRIMARY = 10
90
+ FIRST_ALT = 11
91
+ SECOND_ALT = 12
92
+ THIRD_ALT = 13
93
+ FOURTH_ALT = 14
94
+ FIFTH_ALT = 15
95
+ SIXTH_ALT = 16
96
+ SEVENTH_ALT = 17
97
+ EIGHTH_ALT = 18
98
+ NINTH_ALT = 19
99
+ GOTHIC = 20
100
+ DOUBLE_UNDERLINE = 21
101
+ RESET_BOLD_AND_FAINT = 22
102
+ RESET_ITALIC_AND_GOTHIC = 23
103
+ RESET_UNDERLINES = 24
104
+ RESET_BLINKING = 25
105
+ POSITIVE = 26
106
+ REVEALED_CHARS = 28
107
+ RESET_CROSSED_OUT = 29
108
+ BLACK_FG = 30
109
+ RED_FG = 31
110
+ GREEN_FG = 32
111
+ YELLOW_FG = 33
112
+ BLUE_FG = 34
113
+ MAGENTA_FG = 35
114
+ CYAN_FG = 36
115
+ WHITE_FG = 37
116
+ ANSI_256_SET_FG = 38
117
+ DEFAULT_FG_COLOR = 39
118
+ BLACK_BG = 40
119
+ RED_BG = 41
120
+ GREEN_BG = 42
121
+ YELLOW_BG = 43
122
+ BLUE_BG = 44
123
+ MAGENTA_BG = 45
124
+ CYAN_BG = 46
125
+ WHITE_BG = 47
126
+ ANSI_256_SET_BG = 48
127
+ DEFAULT_BG_COLOR = 49
128
+ FRAMED = 50
129
+ ENCIRCLED = 52
130
+ OVERLINED = 53
131
+ NOT_FRAMED_OR_CIRCLED = 54
132
+ IDEOGRAM_UNDER_OR_RIGHT = 55
133
+ IDEOGRAM_2UNDER_OR_2RIGHT = 60
134
+ IDEOGRAM_OVER_OR_LEFT = 61
135
+ IDEOGRAM_2OVER_OR_2LEFT = 62
136
+ CANCEL = 63
137
+ BLACK_BRIGHT_FG = 90
138
+ RED_BRIGHT_FG = 91
139
+ GREEN_BRIGHT_FG = 92
140
+ YELLOW_BRIGHT_FG = 93
141
+ BLUE_BRIGHT_FG = 94
142
+ MAGENTA_BRIGHT_FG = 95
143
+ CYAN_BRIGHT_FG = 96
144
+ WHITE_BRIGHT_FG = 97
145
+ BLACK_BRIGHT_BG = 100
146
+ RED_BRIGHT_BG = 101
147
+ GREEN_BRIGHT_BG = 102
148
+ YELLOW_BRIGHT_BG = 103
149
+ BLUE_BRIGHT_BG = 104
150
+ MAGENTA_BRIGHT_BG = 105
151
+ CYAN_BRIGHT_BG = 106
152
+ WHITE_BRIGHT_BG = 107
153
+
154
+
155
+ # constant for sgr parameter validation
156
+ _SGR_PARAM_VALUES = frozenset(x.value for x in SgrParameter)
157
+
158
+
159
+ class colorbytes(bytes):
160
+
161
+ @classmethod
162
+ def from_rgb(cls, __rgb):
163
+ """Construct a `colorbytes` object from a dictionary of RGB values.
164
+
165
+ Returns
166
+ -------
167
+ color_bytes : ansicolor4Bit | ansicolor8Bit | ansicolor24Bit
168
+ Constructed from the RGB dictionary, or `__rgb` returned if of same type as `cls`.
169
+
170
+ Raises
171
+ ------
172
+ TypeError
173
+ If `__rgb` is not a dictionary.
174
+
175
+ ValueError
176
+ If an unexpected key or value type is encountered in the RGB dict.
177
+
178
+ Examples
179
+ --------
180
+ >>> rgb_dict = {'fg': (255, 85, 85)}
181
+ >>> old_ansi = ansicolor4Bit.from_rgb(rgb_dict)
182
+ >>> repr(old_ansi)
183
+ "ansicolor4Bit(b'91')"
184
+
185
+ >>> new_ansi = ansicolor24Bit.from_rgb(rgb_dict)
186
+ >>> repr(new_ansi)
187
+ "ansicolor24Bit(b'38;2;255;85;85')"
188
+ """
189
+ if not isinstance(__rgb, Mapping):
190
+ raise TypeError
191
+ if __rgb.keys() not in ({'fg'}, {'bg'}):
192
+ raise ValueError
193
+ rgb = {k: tuple(map(int, v)) if isinstance(v, Iterable) else hex2rgb(v)
194
+ for k, v in __rgb.items()}
195
+
196
+ fmt: AnsiColorType = cls if cls is not colorbytes else DEFAULT_ANSI
197
+ inst = bytes.__new__(fmt, rgb2ansi_color_esc(fmt, *rgb.copy().popitem()))
198
+ setattr(inst, '_rgb_dict_', rgb)
199
+ return cast(AnsiColorFormat, inst)
200
+
201
+ def __new__(cls, __ansi):
202
+ if not isinstance(__ansi, (bytes, bytearray)):
203
+ raise TypeError(
204
+ f"Expected bytes-like object, got {type(__ansi).__name__} instead") from None
205
+ if (is_subtype := cls is not colorbytes) and type(__ansi) is cls:
206
+ return cast(AnsiColorFormat, __ansi)
207
+ match __ansi.removeprefix(CSI).removesuffix(b'm').split(b';'):
208
+ case [color]:
209
+ typ = ansicolor4Bit
210
+ k, rgb = _ANSI16C_I2KV[int(color)]
211
+ case [(b'38' | b'48') as k, b'5', color]:
212
+ typ = ansicolor8Bit
213
+ k = _ANSI256_B2KEY[k]
214
+ rgb = ansi_8bit_to_rgb(int(color))
215
+ case [(b'38' | b'48') as k, b'2', r, g, b]:
216
+ typ = ansicolor24Bit
217
+ k = _ANSI256_B2KEY[k]
218
+ rgb = int(r), int(g), int(b)
219
+ case _:
220
+ raise ValueError
221
+ if typ is not cls:
222
+ __ansi = rgb2ansi_color_esc(
223
+ cls if is_subtype else typ,
224
+ mode=cast(ColorDictKeys, k),
225
+ rgb=rgb)
226
+ inst = bytes.__new__(typ, __ansi)
227
+ setattr(inst, '_rgb_dict_', {k: rgb})
228
+ return cast(AnsiColorFormat, inst)
229
+
230
+ def __repr__(self):
231
+ return f"{type(self).__name__}({super().__repr__()})"
232
+
233
+ @property
234
+ def rgb_dict(self):
235
+ return MappingProxyType(self._rgb_dict_)
236
+
237
+
238
+ class ansicolor4Bit(colorbytes):
239
+ """ANSI 4-bit color format.
240
+
241
+ alias: '4b'
242
+
243
+ Supports 16 colors:
244
+ * 8 standard colors:
245
+ {0: black, 1: red, 2: green, 3: yellow, 4: blue, 5: magenta, 6: cyan, 7: white}
246
+ * 8 bright colors, each mapping to a standard color (bright = standard + 8).
247
+
248
+ Color codes use escape sequences of the form:
249
+ * `CSI 30–37 m` for standard foreground colors.
250
+ * `CSI 40–47 m` for standard background colors.
251
+ * `CSI 90–97 m` for bright foreground colors.
252
+ * `CSI 100–107 m` for bright background colors.
253
+ Where `CSI` (Control Sequence Introducer) is `ESC[`.
254
+
255
+ Examples
256
+ --------
257
+ bright red fg:
258
+ `ESC[91m`
259
+
260
+ standard green bg:
261
+ `ESC[42m`
262
+
263
+ bright white bg, black fg:
264
+ `ESC[107;30m`
265
+ """
266
+ pass
267
+
268
+
269
+ class ansicolor8Bit(colorbytes):
270
+ """ANSI 8-Bit color format.
271
+
272
+ alias: '8b'
273
+
274
+ Supports 256 colors, mapped to the following value ranges:
275
+ * (0, 15): Corresponds to ANSI 4-bit colors.
276
+ * (16, 231): Represents a 6x6x6 RGB color cube.
277
+ * (232, 255): Greyscale colors, from black to white.
278
+
279
+ Color codes use escape sequences of the form:
280
+ * `CSI 38;5;(n) m` for foreground colors.
281
+ * `CSI 48;5;(n) m` for background colors.
282
+ Where `CSI` (Control Sequence Introducer) is `ESC[` and `n` is an unsigned 8-bit integer.
283
+
284
+ Examples
285
+ --------
286
+ white bg:
287
+ `ESC[48;5;255m`
288
+
289
+ bright red fg (ANSI 4-bit):
290
+ `ESC[38;5;9m`
291
+
292
+ bright red fg (color cube):
293
+ `ESC[38;5;196m`
294
+ """
295
+ pass
296
+
297
+
298
+ class ansicolor24Bit(colorbytes):
299
+ """ANSI 24-Bit color format.
300
+
301
+ alias: '24b'
302
+
303
+ Supports all colors in the RGB color space (16,777,216 total).
304
+
305
+ Color codes use escape sequences of the form:
306
+ * `CSI 38;2;(r);(g);(b) m` for foreground colors.
307
+ * `CSI 48;2;(r);(g);(b) m` for background colors.
308
+ Where `CSI` (Control Sequence Introducer) is `ESC[` and `r`, `g`, `b` are unsigned 8-bit ints.
309
+
310
+ Examples
311
+ --------
312
+ red fg:
313
+ `ESC[38;2;255;85;85m`
314
+
315
+ black bg:
316
+ `ESC[48;2;0;0;0m`
317
+
318
+ white fg, green bg:
319
+ `ESC[38;2;255;255;255;48;2;0;170;0m`
320
+ """
321
+ pass
322
+
323
+
324
+ _SUPPORTS_256 = frozenset(
325
+ ['ANSICON',
326
+ 'COLORTERM',
327
+ 'ConEmuANSI',
328
+ 'PYCHARM_HOSTED',
329
+ 'TERM',
330
+ 'TERMINAL_EMULATOR',
331
+ 'TERM_PROGRAM',
332
+ 'WT_SESSION'])
333
+
334
+
335
+ def is_vt_proc_enabled():
336
+ if not _SUPPORTS_256 & os.environ.keys():
337
+ if os.name == 'nt':
338
+ STD_OUTPUT_HANDLE = -11
339
+ ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
340
+ kernel32 = windll.kernel32
341
+ handle = kernel32.GetStdHandle(STD_OUTPUT_HANDLE)
342
+ if handle == -1:
343
+ return False
344
+ mode = c_ulong()
345
+ if not kernel32.GetConsoleMode(handle, byref(mode)):
346
+ return False
347
+ mode.value |= ENABLE_VIRTUAL_TERMINAL_PROCESSING
348
+ if not kernel32.SetConsoleMode(handle, mode):
349
+ return False
350
+ return True
351
+
352
+
353
+ def get_term_ansi_default():
354
+ return ansicolor8Bit if is_vt_proc_enabled() else ansicolor4Bit
355
+
356
+
357
+ DEFAULT_ANSI = get_term_ansi_default()
358
+ _ANSI_COLOR_TYPES = frozenset(colorbytes.__subclasses__())
359
+
360
+
361
+ @lru_cache
362
+ def _is_ansi_type(typ: type):
363
+ try:
364
+ return typ in _ANSI_COLOR_TYPES
365
+ except TypeError:
366
+ return False
367
+
368
+
369
+ AnsiColorFormat = ansicolor4Bit | ansicolor8Bit | ansicolor24Bit
370
+ AnsiColorType = type[AnsiColorFormat]
371
+ AnsiColorParam = AnsiColorAlias | AnsiColorType
372
+ _AnsiColor_co = TypeVar('_AnsiColor_co', bound=colorbytes, covariant=True)
373
+
374
+ _ANSI_FORMAT_MAP = cast(
375
+ dict[AnsiColorParam, AnsiColorType], {
376
+ **dict(x * 2 for x in zip(_ANSI_COLOR_TYPES)),
377
+ **{k.__args__[0]: t
378
+ for k, t in zip(
379
+ sorted(
380
+ AnsiColorAlias.__args__,
381
+ key=lambda x: int(x.__args__[0].removesuffix('b'))),
382
+ sorted(
383
+ _ANSI_COLOR_TYPES,
384
+ key=lambda x: (
385
+ lambda n: int(n[n.index('r') + 1:n.rindex('B')]))(x.__name__)))}
386
+ })
387
+
388
+
389
+ def get_ansi_type(typ):
390
+ try:
391
+ return _ANSI_FORMAT_MAP[typ]
392
+ except (TypeError, KeyError) as e:
393
+ if isinstance(typ, str):
394
+ raise ValueError(
395
+ f"invalid ANSI color format alias: {str(e)}") from None
396
+ repr_getter = lambda t: (t if isinstance(t, type) else type(t)).__name__
397
+ raise TypeError(
398
+ 'Expected {!r} or type[{}], got {!r} instead'.format(
399
+ str.__qualname__,
400
+ ' | '.join(set(map(repr_getter, _ANSI_FORMAT_MAP.values()))),
401
+ repr_getter(typ))) from None
402
+
403
+
404
+ def rgb2ansi_color_esc(ret_format, mode, rgb):
405
+ ret_format = get_ansi_type(ret_format)
406
+ assert len(rgb) == 3, 'length of RGB value is not 3'
407
+ try:
408
+ if ret_format is ansicolor4Bit:
409
+ return b'%d' % _ANSI16C_KV2I[mode, nearest_ansi_4bit_rgb(rgb)]
410
+ return b';'.join(
411
+ b'%d' % b for b in
412
+ [_ANSI256_KEY2I[mode]]
413
+ + ([5, rgb_to_ansi_8bit(rgb)]
414
+ if ret_format is ansicolor8Bit
415
+ else [2, *rgb]))
416
+ except KeyError:
417
+ if isinstance(mode, str):
418
+ raise ValueError(
419
+ f"invalid mode: {mode!r}")
420
+ raise TypeError(
421
+ f"'mode' argument must be {str.__qualname__}, "
422
+ f"not {type(mode).__qualname__}") from None
423
+
424
+
425
+ class Color(int):
426
+
427
+ def __new__(cls, __x):
428
+ """Convert an integer into a `Color` object.
429
+
430
+ Parameters
431
+ ----------
432
+ __x : SupportsInt | Color
433
+ If another Color object is given, immediately return it unchanged.
434
+ Otherwise, the value must be an integer within the range (0, 0xFFFFFF).
435
+
436
+ Returns
437
+ -------
438
+ Color
439
+ A new Color object.
440
+
441
+ Raises
442
+ ------
443
+ TypeError
444
+ If value is of an unexpected type.
445
+ """
446
+ if type(__x) is cls:
447
+ return __x
448
+ if is_hex_rgb(__x, strict=True):
449
+ inst = super().__new__(cls, int(__x))
450
+ inst._rgb_ = hex2rgb(inst)
451
+ return inst
452
+
453
+ def __repr__(self):
454
+ return f"{type(self).__qualname__}({self:#08x})"
455
+
456
+ def __invert__(self):
457
+ return Color(0xFFFFFF ^ self)
458
+
459
+ @classmethod
460
+ def from_rgb(cls, rgb) -> Self:
461
+ inst = super().__new__(cls, rgb2hex(rgb))
462
+ inst._rgb_ = hex2rgb(inst)
463
+ return inst
464
+
465
+ @property
466
+ def rgb(self):
467
+ return getattr(self, '_rgb_')
468
+
469
+
470
+ def randcolor():
471
+ """Return a random color as a :class:`Color` object."""
472
+ return Color.from_bytes(random.randbytes(3))
473
+
474
+
475
+ class SgrParamWrapper:
476
+ __slots__ = '_value_'
477
+
478
+ def __init__(self, value=b''):
479
+ cls, vt = map(type, (self, value))
480
+ if not issubclass(vt, (cls, bytes)):
481
+ raise TypeError(
482
+ f"expected value to be {cls.__qualname__!r} or bytes-like object,"
483
+ f"got {type(value).__qualname__!r} instead")
484
+ self._value_ = value._value_ if vt is cls else value
485
+
486
+ def __hash__(self):
487
+ return hash(self._value_)
488
+
489
+ def __eq__(self, other):
490
+ cls, other_cls = map(type, (self, other))
491
+ if cls is other_cls or issubclass(other_cls, bytes):
492
+ return hash(self) == hash(other)
493
+ return False
494
+
495
+ def __bytes__(self):
496
+ return self._value_.__bytes__()
497
+
498
+ def __repr__(self):
499
+ return f"{type(self).__name__}({self._value_})"
500
+
501
+ def is_same_kind(self, other):
502
+ if self == other:
503
+ return True
504
+ try:
505
+ return next(_iter_sgr(other)) == self._value_
506
+ except (TypeError, StopIteration, RuntimeError):
507
+ return False
508
+
509
+ def is_reset(self):
510
+ return self._value_ == b'0'
511
+
512
+ def is_color(self):
513
+ return isinstance(self._value_, colorbytes)
514
+
515
+
516
+ SgrParamWrapper.__name__ = SgrParameter.__name__.lower()
517
+
518
+
519
+ @lru_cache
520
+ def _get_sgr_bitmask[_T: (bytes, bytearray, Buffer)](__x: _T) -> list[int]:
521
+ """Return a list of integers from a bytestring of ANSI SGR parameters.
522
+
523
+ Bitwise equivalent to `list(map(int, bytes().split(b';')))`.
524
+ """
525
+ __x = __x.removeprefix(CSI)[:idx if ~(idx := __x.find(0x6d)) else None].removesuffix(b'm')
526
+ length = len(__x)
527
+ a, b = map(int.from_bytes, (bytes([0x3b] * length), __x))
528
+ buffer = []
529
+ allocated = []
530
+ alloc = lambda: allocated.append(
531
+ int(''.join(map(str, buffer))))
532
+ prepass = zip(
533
+ map(bool, (~b & a).to_bytes(length=length)),
534
+ (x % 0x30 for x in __x))
535
+ for c, v in prepass:
536
+ if c:
537
+ buffer.append(v)
538
+ else:
539
+ alloc()
540
+ buffer.clear()
541
+ if buffer:
542
+ alloc()
543
+ return allocated
544
+
545
+
546
+ def _iter_normalized_sgr(__iter) -> Iterator[AnsiColorFormat | int]:
547
+ if isinstance(__iter, Buffer):
548
+ yield from _get_sgr_bitmask(__iter)
549
+ else:
550
+ for it in __iter:
551
+ if _is_ansi_type(type(it)):
552
+ yield it
553
+ elif isinstance(it, int):
554
+ yield int(it)
555
+ elif isinstance(it, (Buffer, SgrParamWrapper)):
556
+ if type(it) is SgrParamWrapper:
557
+ it = it._value_
558
+ yield from _get_sgr_bitmask(it)
559
+ else:
560
+ raise TypeError(
561
+ f"Expected {int.__qualname__!r} or bytes-like object, "
562
+ f"got {type(it).__qualname__!r} instead")
563
+
564
+
565
+ def _co_yield_colorbytes(
566
+ __iter: Iterator[int]
567
+ ) -> Generator[bytes | AnsiColorFormat, int, None]:
568
+ m: dict[int, ColorDictKeys] = {38: 'fg', 48: 'bg'}
569
+ key_pair = m.get
570
+ get_4b = _ANSI16C_I2KV.get
571
+ new_4b = lambda t: ansicolor4Bit.from_rgb({t[0]: t[1]})
572
+ new_8b = lambda *args: ansicolor8Bit(
573
+ b';'.join(map(b'%d'.__mod__, (args[0], args[1], next(__iter)))))
574
+ new_24b = lambda x: ansicolor24Bit.from_rgb({x: tuple(next(__iter) for _ in range(3))})
575
+ default = lambda x: bytes(ascii(x), 'ansi')
576
+ obj = bytes()
577
+ while True:
578
+ value = yield obj
579
+ if key := key_pair(value):
580
+ kind = next(__iter)
581
+ if kind == 5:
582
+ obj = new_8b(value, kind)
583
+ else:
584
+ obj = new_24b(key)
585
+ elif kv := get_4b(value):
586
+ obj = new_4b(kv)
587
+ else:
588
+ obj = default(value)
589
+
590
+
591
+ def _gen_colorbytes(__iter: Iterable[int]) -> Iterator[bytes | AnsiColorFormat]:
592
+ gen = iter(__iter)
593
+ color_coro = _co_yield_colorbytes(gen)
594
+ next(color_coro)
595
+ while True:
596
+ try:
597
+ value = next(gen)
598
+ if _is_ansi_type(type(value)):
599
+ yield value
600
+ continue
601
+ yield color_coro.send(value)
602
+ except StopIteration:
603
+ break
604
+
605
+
606
+ def _iter_sgr(__x):
607
+ if isinstance(__x, int):
608
+ __x = [__x]
609
+ return _gen_colorbytes(_iter_normalized_sgr(__x))
610
+
611
+
612
+ class SgrSequence:
613
+
614
+ def append(self, __value):
615
+ if __value not in _SGR_PARAM_VALUES:
616
+ raise ValueError(
617
+ f"{__value!r} is not a valid SGR parameter")
618
+ if kv := _ANSI16C_I2KV.get(__value):
619
+ if __value in _ANSI16C_BRIGHT:
620
+ self._has_bright_colors_ = True
621
+ elif ~(b_idx := self.find(b'1')):
622
+ self._has_bright_colors_ = True
623
+ self.pop(b_idx)
624
+ kv = _ANSI16C_I2KV.get(__value + 60)
625
+ else:
626
+ self._has_bright_colors_ = False
627
+ value = ansicolor4Bit.from_rgb({kv[0]: kv[1]})
628
+ else:
629
+ value = b'%d' % __value
630
+ v = SgrParamWrapper(value)
631
+ if v.is_color():
632
+ for key in (rgb_dict := v._value_._rgb_dict_).keys() & self._rgb_dict_:
633
+ color = self.get_color(key)
634
+ self.pop(self.index(color))
635
+ self._rgb_dict_ |= rgb_dict
636
+ self._sgr_params_.append(v)
637
+ self._bytes_ = None
638
+
639
+ def find(self, value):
640
+ try:
641
+ return self.index(value)
642
+ except ValueError:
643
+ return -1
644
+
645
+ def get_color(self, __key: ColorDictKeys):
646
+ if self.is_color():
647
+ return next((v for v in self if v.is_color() and __key in v._value_.rgb_dict), None)
648
+ return None
649
+
650
+ def index(self, value):
651
+ try:
652
+ return next(i for i, p in enumerate(self) if p.is_same_kind(value))
653
+ except StopIteration:
654
+ raise ValueError(
655
+ f"{value!r} not in sequence") from None
656
+
657
+ def is_color(self):
658
+ return any(p.is_color() for p in self)
659
+
660
+ def is_reset(self):
661
+ return any(p.is_reset() for p in self)
662
+
663
+ def remove(self, value):
664
+ try:
665
+ self.pop(self.index(value))
666
+ except ValueError as e:
667
+ raise ValueError(
668
+ e) from None
669
+
670
+ def pop(self, __index=-1):
671
+ try:
672
+ obj = self._sgr_params_.pop(__index)
673
+ except IndexError as e:
674
+ raise IndexError(
675
+ e) from None
676
+ v = obj._value_
677
+ if obj.is_color():
678
+ for k in v._rgb_dict_.keys():
679
+ del self._rgb_dict_[k]
680
+ if self.has_bright_colors and (vx := int(v)) in _ANSI16C_I2KV:
681
+ self._has_bright_colors_ = False
682
+ if vx in _ANSI16C_STD:
683
+ self.pop(self.index(b'1'))
684
+ elif self.has_bright_colors and v == b'1':
685
+ for p in self._sgr_params_:
686
+ if type(px := p._value_) is not ansicolor4Bit or int(px) not in _ANSI16C_STD:
687
+ continue
688
+ self._has_bright_colors_ = False
689
+ break
690
+ self._bytes_ = None
691
+ return obj
692
+
693
+ def values(self):
694
+ return [p._value_ for p in self._sgr_params_]
695
+
696
+ def __add__(self, other):
697
+ if type(self) is type(other):
698
+ return SgrSequence([*self, *other])
699
+ if isinstance(other, str):
700
+ return str(self) + other
701
+ raise TypeError(
702
+ f"can only concatenate {SgrSequence.__qualname__} "
703
+ f"(not {type(other).__qualname__!r}) to {SgrSequence.__qualname__}")
704
+
705
+ def __bool__(self):
706
+ return bool(self._sgr_params_)
707
+
708
+ def __bytes__(self):
709
+ if self._bytes_ is None:
710
+ if self._sgr_params_:
711
+ self._bytes_ = b'\x1b[%sm' % b';'.join(self.values())
712
+ else:
713
+ self._bytes_ = bytes()
714
+ return self._bytes_
715
+
716
+ def __contains__(self, item: ...):
717
+ if self:
718
+ try:
719
+ return set(_iter_sgr(item)).issubset(self.values())
720
+ except (TypeError, RuntimeError):
721
+ pass
722
+ return False
723
+
724
+ def __copy__(self):
725
+ cls = type(self)
726
+ inst = object.__new__(cls)
727
+ inst._bytes_ = self._bytes_
728
+ inst._has_bright_colors_ = self._has_bright_colors_
729
+ inst._sgr_params_ = self._sgr_params_.copy()
730
+ inst._rgb_dict_ = self._rgb_dict_.copy()
731
+ return inst
732
+
733
+ def __deepcopy__(self, memo):
734
+ cls = type(self)
735
+ inst = object.__new__(cls)
736
+ memo[id(self)] = inst
737
+ inst._bytes_ = self._bytes_
738
+ inst._has_bright_colors_ = self._has_bright_colors_
739
+ inst._sgr_params_ = deepcopy(self._sgr_params_, memo)
740
+ inst._rgb_dict_ = deepcopy(self._rgb_dict_, memo)
741
+ return inst
742
+
743
+ def __eq__(self, other: ...):
744
+ if type(self) is type(other):
745
+ other: SgrSequence
746
+ try:
747
+ return all(
748
+ self_param == other_param for self_param, other_param in
749
+ zip(self.values(), other.values(), strict=True))
750
+ except ValueError:
751
+ return False
752
+ return False
753
+
754
+ def __getitem__(self, item):
755
+ return self._sgr_params_[item]
756
+
757
+ def __iadd__(self, other: 'SgrSequence'):
758
+ if type(self) is not type(other):
759
+ raise TypeError(
760
+ f"can only concatenate {SgrSequence.__qualname__} "
761
+ f"(not {type(other).__qualname__!r}) to {SgrSequence.__qualname__}")
762
+ return SgrSequence(self._sgr_params_ + other._sgr_params_)
763
+
764
+ def __init__(self, __iter=None, *, ansi_type=None) -> None:
765
+ cls = type(self)
766
+ if type(__iter) is cls:
767
+ other = __iter.__copy__()
768
+ for attr in cls.__slots__:
769
+ setattr(self, attr, getattr(other, attr))
770
+ return
771
+
772
+ self._bytes_ = None
773
+ self._has_bright_colors_ = False
774
+ self._rgb_dict_ = {}
775
+ self._sgr_params_ = []
776
+
777
+ if not __iter:
778
+ return
779
+
780
+ values = set()
781
+ add_unique = values.add
782
+ append_param = self._sgr_params_.append
783
+ remove_param = self._sgr_params_.remove
784
+ fg_slot: SgrParamWrapper | None
785
+ bg_slot: SgrParamWrapper | None
786
+ color_dict = dict.fromkeys(['fg', 'bg'], None)
787
+ is_bold = has_bold = False
788
+
789
+ def update_colors(
790
+ __param: SgrParamWrapper,
791
+ __rgb_dict: Mapping[ColorDictKeys, Int3Tuple]
792
+ ):
793
+ k: ColorDictKeys
794
+ for k, slot in color_dict.items():
795
+ if v := __rgb_dict.get(k):
796
+ if slot:
797
+ remove_param(slot)
798
+ color_dict[k] = __param
799
+ self._rgb_dict_[k] = v
800
+
801
+ is_diff_ansi_typ: Callable[[AnsiColorFormat], bool]
802
+ if ansi_type is None:
803
+ is_diff_ansi_typ = lambda _: False
804
+ else:
805
+ assert ansi_type in _ANSI_COLOR_TYPES
806
+ is_diff_ansi_typ = lambda v: type(v) is not ansi_type
807
+
808
+ for x in _iter_sgr(__iter):
809
+ if x in values:
810
+ continue
811
+ param = SgrParamWrapper(x)
812
+ if x == b'1':
813
+ if not is_bold:
814
+ has_bold = is_bold = True
815
+ elif hasattr(x, 'rgb_dict'):
816
+ if is_diff_ansi_typ(x):
817
+ param = SgrParamWrapper(x := ansi_type.from_rgb(x))
818
+ if type(x) is ansicolor4Bit:
819
+ if (btoi := int(x)) in _ANSI16C_BRIGHT:
820
+ self._has_bright_colors_ = True
821
+ elif is_bold and btoi in _ANSI16C_STD:
822
+ self._has_bright_colors_ = True
823
+ param = SgrParamWrapper(x := ansicolor4Bit(b'%d' % (btoi + 60)))
824
+ if has_bold:
825
+ self._sgr_params_.pop(
826
+ next(
827
+ i for i, v in enumerate(self._sgr_params_)
828
+ if v._value_ == b'1'))
829
+ has_bold = False
830
+ update_colors(param, x.rgb_dict)
831
+ append_param(param)
832
+ add_unique(x)
833
+
834
+ if self._sgr_params_[-1]._value_ == b'0':
835
+ self._has_bright_colors_ = False
836
+ self._sgr_params_ = [self._sgr_params_.pop()]
837
+ self._rgb_dict_ = {}
838
+ self._bytes_ = b'\x1b[%sm' % b';'.join(map(bytes, self._sgr_params_))
839
+
840
+ def __iter__(self):
841
+ return iter(self._sgr_params_)
842
+
843
+ def __radd__(self, other):
844
+ if type(self) is type(other):
845
+ return SgrSequence([*other, *self])
846
+ if isinstance(other, str):
847
+ return other + str(self)
848
+ raise TypeError(
849
+ f"can only concatenate {SgrSequence.__qualname__} "
850
+ f"(not {type(other).__qualname__!r}) to {SgrSequence.__qualname__}")
851
+
852
+ def __repr__(self):
853
+ return f"{type(self).__qualname__}({self.values()})"
854
+
855
+ def __str__(self):
856
+ return str(bytes(self), 'utf-8')
857
+
858
+ __slots__ = '_bytes_', '_has_bright_colors_', '_rgb_dict_', '_sgr_params_'
859
+
860
+ @property
861
+ def bg(self):
862
+ return self.rgb_dict.get('bg')
863
+
864
+ @property
865
+ def fg(self):
866
+ return self.rgb_dict.get('fg')
867
+
868
+ @property
869
+ def has_bright_colors(self):
870
+ return self._has_bright_colors_
871
+
872
+ @property
873
+ def rgb_dict(self):
874
+ return MappingProxyType(self._rgb_dict_)
875
+
876
+ @rgb_dict.deleter
877
+ def rgb_dict(self) -> None:
878
+ for k in self._rgb_dict_.keys():
879
+ self.pop(self.index(self.get_color(k)))
880
+ self._bytes_ = None
881
+
882
+ @rgb_dict.setter
883
+ def rgb_dict[_AnsiColorType: type[AnsiColorFormat]](
884
+ self,
885
+ __value: tuple[_AnsiColorType, dict[ColorDictKeys, Union[Color, None]]]
886
+ ) -> None:
887
+ ansi_type, color_dict = __value
888
+ for k, v in color_dict.items():
889
+ if v is not None:
890
+ if self._rgb_dict_.get(k):
891
+ try:
892
+ self.pop(self.index(self.get_color(k)))
893
+ except ValueError as e:
894
+ e.add_note(repr(self))
895
+ raise e
896
+ color_bytes = ansi_type.from_rgb({k: v})
897
+ self._rgb_dict_ |= color_bytes._rgb_dict_
898
+ self._sgr_params_.append(SgrParamWrapper(color_bytes))
899
+ self._bytes_ = None
900
+
901
+
902
+ # type alias types for ColorStr constructor `color_spec` parameter forms
903
+ type _CSpecScalar = Union[int, Color, RGBVectorLike]
904
+ type _CSpecKVPair = tuple[ColorDictKeys, _CSpecScalar]
905
+ type _CSpecTuplePair = tuple[_CSpecScalar, _CSpecScalar] | tuple[_CSpecKVPair, _CSpecKVPair]
906
+ type _CSpecDict = Mapping[ColorDictKeys, _CSpecScalar]
907
+ type _CSpecType = (Union[SgrSequence, str, bytes]
908
+ | _CSpecScalar
909
+ | _CSpecTuplePair
910
+ | _CSpecKVPair
911
+ | _CSpecDict)
912
+
913
+ _ColorSpec = TypeVar('_ColorSpec', bound=_CSpecType)
914
+
915
+
916
+ class _ColorDict(TypedDict, total=False):
917
+ fg: Union[Color, AnsiColorFormat, None]
918
+ bg: Union[Color, AnsiColorFormat, None]
919
+
920
+
921
+ def _solve_color_spec[_T: (_CSpecType, SgrSequence)](
922
+ color_spec: _T | None,
923
+ ansi_type: type[AnsiColorFormat]
924
+ ):
925
+ keys = ['bg', 'fg']
926
+ valid_keys = set(keys)
927
+
928
+ def resolve(value, *, key=None):
929
+ nonlocal keys
930
+ if key is not None:
931
+ assert key in valid_keys, 'expected literal keys {}, got {!r}'.format(valid_keys, key)
932
+ if key in keys:
933
+ keys.remove(key)
934
+ match value:
935
+ case Color() | int() | np.integer() as color:
936
+ yield (key or keys.pop(), Color(color).rgb)
937
+ case [int(), int(), int()] as rgb:
938
+ r, g, b = (x & 0xff for x in rgb)
939
+ yield (key or keys.pop(), (r, g, b))
940
+ case np.ndarray() as colors:
941
+ assert not colors.shape[-1] % 3, 'array does not contain RGB values'
942
+ it = np.uint8(colors).flat
943
+ for _ in range(colors.ndim):
944
+ yield (key or keys.pop(), tuple(int(next(it)) for _ in range(3)))
945
+ case {'fg': _, 'bg': _} | {'fg': _} | {'bg': _} as colors:
946
+ for key, color in colors.items():
947
+ yield from resolve(color, key=key)
948
+ case [str() as key, color] if key in valid_keys:
949
+ yield from resolve(color, key=key)
950
+ case [_, _] as colors:
951
+ for color in colors:
952
+ yield from resolve(color)
953
+ case _:
954
+ raise ValueError(repr(value))
955
+
956
+ out = dict()
957
+ try:
958
+ for k, v in resolve(color_spec):
959
+ if k in out:
960
+ if len(out) > 1 and out[k] != v:
961
+ raise ValueError(
962
+ f"multiple possible values for {k!r} {(out[k], v)}")
963
+ out[keys.pop()] = out.pop(k)
964
+ keys.append(k)
965
+ out[k] = v
966
+ except Exception as e:
967
+ if type(e) is IndexError:
968
+ e = ValueError(
969
+ 'too many arguments'
970
+ if len(out) >= 2 else
971
+ 'args contain non-RGB values')
972
+ context = ('invalid color spec', str(e))
973
+ raise ValueError(
974
+ ': '.join(filter(None, context))) from None
975
+ return SgrSequence([ansi_type.from_rgb({k: v}) for k, v in out.items()])
976
+
977
+
978
+ def _get_color_str_vars(base_str: Optional[str],
979
+ color_spec: Optional[_ColorSpec],
980
+ ansi_type: AnsiColorType = None) -> tuple[SgrSequence, str]:
981
+ if color_spec is None:
982
+ return SgrSequence(), base_str or ''
983
+ if ansi_type is None:
984
+ ansi_type = DEFAULT_ANSI
985
+ if isinstance(color_spec, (str, bytes)):
986
+ if hasattr(color_spec, 'encode'):
987
+ color_spec = color_spec.encode()
988
+ if csi_count := color_spec.count(CSI):
989
+ if csi_count > 1:
990
+ color_spec, _, byte_str = (
991
+ color_spec
992
+ .removeprefix(CSI)
993
+ .removesuffix(SGR_RESET.encode())
994
+ .partition(b'm')
995
+ )
996
+ if color_spec.count(CSI) > 1:
997
+ raise ValueError(
998
+ f"color spec contains {csi_count} escape sequences, expected only 1"
999
+ ) from None
1000
+ base_str = byte_str.decode()
1001
+ sgr_params = SgrSequence(color_spec, ansi_type=ansi_type)
1002
+ else:
1003
+ is_hex_rgb(color_spec := int.from_bytes(color_spec), strict=True)
1004
+ sgr_params = _solve_color_spec(color_spec, ansi_type=ansi_type)
1005
+ elif not isinstance(color_spec, SgrSequence):
1006
+ sgr_params = _solve_color_spec(color_spec, ansi_type=ansi_type)
1007
+ else:
1008
+ sgr_params = color_spec
1009
+ base_str = base_str or ''
1010
+ return sgr_params, base_str
1011
+
1012
+
1013
+ class _ColorStrWeakVars(TypedDict, total=False):
1014
+ _base_str_: str
1015
+ _sgr_: SgrSequence
1016
+ _no_reset_: bool
1017
+
1018
+
1019
+ class _AnsiBytesGetter:
1020
+
1021
+ def __get__(self, instance: Union['ColorStr', None], objtype=None):
1022
+ if instance is None:
1023
+ return
1024
+ return instance._sgr_.__bytes__()
1025
+
1026
+
1027
+ class _SgrParamsGetter:
1028
+
1029
+ def __get__(self, instance: Union['ColorStr', None], objtype=None):
1030
+ if instance is None:
1031
+ return
1032
+ return instance._sgr_._sgr_params_
1033
+
1034
+
1035
+ class _ColorDictGetter:
1036
+
1037
+ def __get__(self, instance: Union['ColorStr', None], objtype=None):
1038
+ if instance is None:
1039
+ return
1040
+ return {k: Color.from_rgb(v) for k, v in instance._sgr_.rgb_dict.items()}
1041
+
1042
+
1043
+ class ColorStr(str):
1044
+
1045
+ def _weak_var_update(self, **kwargs):
1046
+ if kwargs.keys().isdisjoint(inst_vars := vars(self)):
1047
+ raise ValueError(
1048
+ f"unexpected keys: {kwargs.keys() - inst_vars.keys()}"
1049
+ ) from None
1050
+ sgr = kwargs.get('_sgr_', self._sgr_)
1051
+ base_str = kwargs.get('_base_str_', self._base_str_)
1052
+ suffix = '' if kwargs.get('_no_reset_', self._no_reset_) else SGR_RESET
1053
+ inst = super().__new__(ColorStr, ''.join([str(sgr), base_str, suffix]))
1054
+ inst.__dict__ |= {**inst_vars, **kwargs}
1055
+ return cast(ColorStr, inst)
1056
+
1057
+ def ansi_partition(self):
1058
+ """Returns the 3-tuple: SGR sequence prefix, base string, SGR reset (or empty string)."""
1059
+ return str(self._sgr_), self.base_str, '' if self.no_reset else SGR_RESET
1060
+
1061
+ def as_ansi_type(self, __ansi_type):
1062
+ """Convert all ANSI color codes in the :class:`ColorStr` to a single ANSI type.
1063
+
1064
+ Parameters
1065
+ ----------
1066
+ __ansi_type : str or type[ansicolor4Bit | ansicolor8Bit | ansicolor24Bit]
1067
+ ANSI format to which all SGR parameters of type :class:`colorbytes` will be cast.
1068
+
1069
+ Returns
1070
+ -------
1071
+ ColorStr
1072
+ Return `self` if all ANSI formats are already the input type.
1073
+ Otherwise, return reformatted :class:`ColorStr`.
1074
+
1075
+ """
1076
+ ansi_type = get_ansi_type(__ansi_type)
1077
+ if self._sgr_.is_color():
1078
+ new_params = []
1079
+ new_rgb = {}
1080
+ for p in self._sgr_params_:
1081
+ if p.is_color() and type(p._value_) is not ansi_type:
1082
+ new_ansi = ansi_type(p._value_)
1083
+ new_rgb |= new_ansi.rgb_dict
1084
+ new_params.append(SgrParamWrapper(new_ansi))
1085
+ else:
1086
+ new_params.append(p)
1087
+ if new_params == self._sgr_params_:
1088
+ return self
1089
+ new_sgr = SgrSequence()
1090
+ for name, value in zip(
1091
+ ('_sgr_params_', '_rgb_dict_'),
1092
+ (new_params, new_rgb)):
1093
+ setattr(new_sgr, name, value)
1094
+ inst = super().__new__(
1095
+ type(self),
1096
+ ''.join([str(new_sgr), self._base_str_, '' if self._no_reset_ else SGR_RESET]))
1097
+ for name, value in {
1098
+ **vars(self),
1099
+ '_sgr_': new_sgr,
1100
+ '_ansi_type_': ansi_type
1101
+ }.items():
1102
+ setattr(inst, name, value)
1103
+ return cast(ColorStr, inst)
1104
+ return self
1105
+
1106
+ def format(self, *args, **kwargs):
1107
+ return self._weak_var_update(_base_str_=self.base_str.format(*args, **kwargs))
1108
+
1109
+ def split(self, sep=None, maxsplit=-1):
1110
+ return list(
1111
+ self._weak_var_update(_base_str_=s) for s in
1112
+ self.base_str.split(sep=sep, maxsplit=maxsplit))
1113
+
1114
+ def recolor(self, __value=None, absolute=False, **kwargs):
1115
+ """Return a copy of `self` with a new color spec.
1116
+
1117
+ If `__value` is a :class:`ColorStr`, return `self` with the colors of `__value`.
1118
+
1119
+ Parameters
1120
+ ----------
1121
+ __value : ColorStr, optional
1122
+ A :class:`ColorStr` object that the new instance will inherit colors from.
1123
+
1124
+ absolute : bool
1125
+ If True, overwrite all colors of the current object with the provided arguments,
1126
+ removing any existing colors not explicitly set by the arguments.
1127
+ Otherwise, only replace colors where specified (default).
1128
+
1129
+ Keyword Args
1130
+ ------------
1131
+ fg : Color, optional
1132
+ New foreground color.
1133
+
1134
+ bg : Color, optional
1135
+ New background color.
1136
+
1137
+ Returns
1138
+ -------
1139
+ ColorStr
1140
+ A new :class:`ColorStr` instance recolored by the input parameters.
1141
+
1142
+ Raises
1143
+ ------
1144
+ TypeError
1145
+ If `__value` is not None but is not an instance of :class:`ColorStr`.
1146
+
1147
+ ValueError
1148
+ If any unexpected keys or value types found in `kwargs`.
1149
+
1150
+ Examples
1151
+ --------
1152
+ >>> cs1 = ColorStr('foo', randcolor())
1153
+ >>> cs2 = ColorStr('bar', dict(fg=Color(0xFF5555), bg=Color(0xFF00FF)))
1154
+ >>> new_cs = cs2.recolor(bg=cs1.fg)
1155
+ >>> int(new_cs.fg) == 0xFF5555, new_cs.bg == cs1.fg
1156
+ (True, True)
1157
+
1158
+ >>> cs = ColorStr("Red text", ('fg', 0xFF0000))
1159
+ >>> recolored = cs.recolor(fg=Color(0x00FF00))
1160
+ >>> recolored.base_str, f"{recolored.fg:06X}"
1161
+ ('Red text', '00FF00')
1162
+ """
1163
+ if __value:
1164
+ if isinstance(__value, ColorStr):
1165
+ kwargs = __value._color_dict_
1166
+ else:
1167
+ raise TypeError(
1168
+ f"expected positional argument of type {ColorStr.__qualname__!r}, "
1169
+ f"got {type(__value).__qualname__!r} instead") from None
1170
+ elif not kwargs:
1171
+ return self
1172
+ valid, context = is_matching_typed_dict(kwargs, _ColorDict)
1173
+ if not valid:
1174
+ raise ValueError(
1175
+ context)
1176
+ sgr = SgrSequence(self._sgr_)
1177
+ if bool(absolute):
1178
+ del sgr.rgb_dict
1179
+ sgr.rgb_dict = (self.ansi_format, kwargs)
1180
+ return self._weak_var_update(_sgr_=sgr)
1181
+
1182
+ def replace(self, __old, __new, __count=-1):
1183
+ if isinstance(__new, ColorStr):
1184
+ __new = __new.base_str
1185
+ return self._weak_var_update(_base_str_=self.base_str.replace(__old, __new, __count))
1186
+
1187
+ def translate(self, __table) -> 'ColorStr':
1188
+ return self._weak_var_update(_base_str_=self.base_str.translate(__table))
1189
+
1190
+ def update_sgr(self, *p):
1191
+ """Return a copy of `self` with updated SGR sequence parameters.
1192
+
1193
+ Parameters
1194
+ ----------
1195
+ *p: SgrParameter | int
1196
+ The SGR parameter value(s) to be added or removed from the :class:`ColorStr`.
1197
+ A value already in `self` SGR sequence gets removed, else it gets added.
1198
+ If no values are passed, returns `self` unchanged.
1199
+
1200
+ Returns
1201
+ -------
1202
+ ColorStr
1203
+ A new :class:`ColorStr` object with the SGR updates applied.
1204
+
1205
+ Raises
1206
+ ------
1207
+ ValueError
1208
+ If any of the SGR parameters are invalid, or if extended color codes are passed.
1209
+
1210
+ Notes
1211
+ -----
1212
+ * The extended color escapes `{38, 48}` require extra parameters and so raise a ValueError.
1213
+ :meth:`ColorStr.as_ansi_type` should be used to change ANSI color format instead.
1214
+
1215
+ Examples
1216
+ --------
1217
+ >>> # creating an empty ColorStr object
1218
+ >>> empty_cs = ColorStr(no_reset=True)
1219
+ >>> empty_cs.ansi
1220
+ b''
1221
+
1222
+ >>> # adding red foreground color
1223
+ >>> red_fg = empty_cs.update_sgr(SgrParameter.RED_FG)
1224
+ >>> red_fg.rgb_dict
1225
+ {'fg': (170, 0, 0)}
1226
+
1227
+ >>> # removing the same parameter
1228
+ >>> empty_cs = red_fg.update_sgr(31)
1229
+ >>> empty_cs.ansi, empty_cs.rgb_dict
1230
+ (b'', {})
1231
+
1232
+ >>> # adding more parameters
1233
+ >>> styles = [SgrParameter.BOLD, SgrParameter.ITALICS, SgrParameter.NEGATIVE]
1234
+ >>> stylized_cs = empty_cs.update_sgr(*styles)
1235
+ >>> stylized_cs.ansi.replace(CSI, b'ESC[')
1236
+ b'ESC[1;3;7m'
1237
+
1238
+ >>> # parameter updates also supported by the `__add__` operator
1239
+ >>> stylized_cs += SgrParameter.BLACK_BG # add background color
1240
+ >>> stylized_cs += SgrParameter.BOLD # remove bold style
1241
+ >>> stylized_cs.ansi.replace(CSI, b'ESC['), stylized_cs.rgb_dict
1242
+ (b'ESC[3;7;40m', {'bg': (0, 0, 0)})
1243
+ """
1244
+ if not p:
1245
+ return self
1246
+ if any(
1247
+ not isinstance(x, int)
1248
+ or x in _ANSI256_KEY2I.values()
1249
+ for x in p):
1250
+ raise ValueError
1251
+ new_sgr = SgrSequence(self._sgr_)
1252
+ for x in p:
1253
+ if x in new_sgr:
1254
+ new_sgr.pop(new_sgr.index(x))
1255
+ elif x == 1 and new_sgr.has_bright_colors:
1256
+ for i, param in enumerate(new_sgr):
1257
+ if type(px := param._value_) is not ansicolor4Bit:
1258
+ continue
1259
+ new_sgr.pop(i)
1260
+ new_sgr.append(int(px) - 60)
1261
+ else:
1262
+ new_sgr.append(x)
1263
+ if new_sgr.is_color():
1264
+ formats: list[AnsiColorType] = [type(p._value_) for p in new_sgr if p.is_color()]
1265
+ ansi_type = max(formats, key=formats.count)
1266
+ else:
1267
+ ansi_type = self.ansi_format
1268
+ inst = super().__new__(
1269
+ type(self),
1270
+ ''.join([str(new_sgr), self._base_str_, '' if self._no_reset_ else SGR_RESET]))
1271
+ inst.__dict__ |= {
1272
+ **vars(self),
1273
+ '_sgr_': new_sgr,
1274
+ '_ansi_type_': ansi_type}
1275
+ return cast(ColorStr, inst)
1276
+
1277
+ def __add__(self, other):
1278
+ if type(self) is type(other):
1279
+ return self._weak_var_update(
1280
+ _sgr_=self._sgr_ + other._sgr_,
1281
+ _base_str_=''.join([self._base_str_, other._base_str_]))
1282
+ if isinstance(other, str):
1283
+ return self._weak_var_update(
1284
+ _base_str_=''.join([self._base_str_, other]))
1285
+ if isinstance(other, SgrParameter):
1286
+ return self.update_sgr(other)
1287
+ if hasattr(other, '_sgr_'):
1288
+ return NotImplemented
1289
+ raise TypeError(
1290
+ f"can only concatenate "
1291
+ f"{str.__name__}, {ColorStr.__name__}, or {SgrParameter.__name__} "
1292
+ f"(got {type(other).__qualname__!r}) "
1293
+ f"to {type(self).__name__}")
1294
+
1295
+ def __contains__(self, __key: str):
1296
+ if type(__key) is not str:
1297
+ return False
1298
+ if __key == str(self._sgr_):
1299
+ return True
1300
+ if __key == SGR_RESET:
1301
+ return not self.no_reset
1302
+ return self.base_str.__contains__(__key)
1303
+
1304
+ def __eq__(self, other):
1305
+ if type(self) is type(other):
1306
+ return hash(self) == hash(other)
1307
+ return False
1308
+
1309
+ def __format__(self, format_spec=''):
1310
+ if ansi_typ := {
1311
+ '4b': ansicolor4Bit,
1312
+ '8b': ansicolor8Bit,
1313
+ '24b': ansicolor24Bit
1314
+ }.get(format_spec):
1315
+ return str(self.as_ansi_type(ansi_typ))
1316
+ return str.__format__(self, format_spec)
1317
+
1318
+ def __getitem__(self, __key: Union[SupportsIndex, slice]):
1319
+ return self._weak_var_update(_base_str_=self.base_str[__key])
1320
+
1321
+ def __hash__(self):
1322
+ return str(self).__hash__()
1323
+
1324
+ # noinspection PyUnusedLocal
1325
+ def __init__(self, obj=None, color_spec=None, **kwargs):
1326
+ """
1327
+ Create a ColorStr object.
1328
+
1329
+ Parameters
1330
+ ----------
1331
+ obj : object, optional
1332
+ The base object to be cast to a ColorStr. If None, uses a null string ('').
1333
+
1334
+ color_spec : type[_ColorSpec | ColorStr], optional
1335
+ The color specification for the string.
1336
+ The constructor supports various types, such as:
1337
+
1338
+ * An RGB tuple
1339
+ * A hex color as an integer
1340
+ * A Color object
1341
+ * Any tuple pair of the aforementioned types:
1342
+ ('fg'=color_spec[0], 'bg'=color_spec[1])
1343
+ * A key-value pair or `dict_items`-like tuple:
1344
+ ('fg', ...) or (('fg', ...), ('bg', ...))
1345
+ * A dictionary mapping:
1346
+ dict[Literal['fg', 'bg'], ...]
1347
+
1348
+ Keyword Args
1349
+ ------------
1350
+ ansi_type : str or type[ansicolor4Bit | ansicolor8Bit | ansicolor24Bit], optional
1351
+ An ANSI format to cast all :class:`colorbytes` params to before formatting the string.
1352
+
1353
+ * ANSI format can also be changed on instances using :meth:`ColorStr.as_ansi_type`
1354
+ * Reformatting recursively applies to `alt_spec` if `alt_spec` is not None
1355
+
1356
+ no_reset : bool
1357
+ If True, create the :class:`ColorStr` without concatenating a 'reset all' SGR sequence.
1358
+ Default is False (new instances get concatenated with reset sequences).
1359
+
1360
+ Returns
1361
+ -------
1362
+ ColorStr
1363
+ A new ColorStr object comprised of the base string and provided ANSI sequences.
1364
+
1365
+ Notes
1366
+ -----
1367
+ * Each of the ANSI color formats can be invoked by their alias in place of the type:
1368
+ ``ansicolor4Bit`` == '4b', ``ansicolor8Bit`` == '8b', ``ansicolor24Bit`` == '24b'
1369
+ * Use :py:func:`help` with :class:`colorbytes` types for color code ranges and sequences.
1370
+
1371
+ * ``color_spec`` of type :obj:`str` or :obj:`bytes` is parsed as a literal escape sequence.
1372
+
1373
+ Examples
1374
+ --------
1375
+ >>> cs = ColorStr('Red text', ('fg', 0xFF0000))
1376
+ >>> cs.rgb_dict, cs.base_str
1377
+ ({'fg': (255, 0, 0)}, 'Red text')
1378
+
1379
+ >>> cs_from_rgb = ColorStr(color_spec={'fg': (255, 85, 85)}, ansi_type='4b')
1380
+ >>> cs_from_literal = ColorStr(color_spec='\x1b[91m', ansi_type='4b')
1381
+ >>> cs_from_rgb == cs_from_literal
1382
+ True
1383
+
1384
+ >>> # ANSI 4-bit sequences of the form `ESC[<1 (bold)>;<{30-37} | {40-47}>...`
1385
+ >>> # are equivalent to 'bright' counterparts `ESC[<{90-97} | {100-107}>...`
1386
+ >>> cs_from_literal_alt = ColorStr(color_spec='\x1b[1;31m', ansi_type='4b')
1387
+ >>> cs_from_literal_alt == cs_from_literal
1388
+ True
1389
+
1390
+ >>> # bold-prefix syntax is autocast to the 'bright' sequence form
1391
+ >>> cs_from_literal_alt.ansi.replace(CSI, b'ESC[')
1392
+ b'ESC[91m'
1393
+ """
1394
+ ...
1395
+
1396
+ def __iter__(self):
1397
+ yield from map(lambda c: self._weak_var_update(_base_str_=c), self.base_str)
1398
+
1399
+ def __len__(self):
1400
+ return self.base_str.__len__()
1401
+
1402
+ def __matmul__(self, other):
1403
+ """Return a new :class:`ColorStr` with the base string of `self` and colors of `other`"""
1404
+ if type(self) is type(other):
1405
+ return self._weak_var_update(_sgr_=other._sgr_, _no_reset_=other.no_reset)
1406
+ raise TypeError(
1407
+ 'unsupported operand type(s) for @: '
1408
+ f"{type(self).__qualname__!r} and {type(other).__qualname__!r}")
1409
+
1410
+ def __mod__(self, __value):
1411
+ return self._weak_var_update(_base_str_=self.base_str.__mod__(__value))
1412
+
1413
+ def __mul__(self, __value):
1414
+ return self._weak_var_update(_base_str_=self.base_str.__mul__(__value))
1415
+
1416
+ def __invert__(self):
1417
+ """Return a copy of `self` with inverted colors (XORed by '0xFFFFFF')"""
1418
+ sgr = SgrSequence(self._sgr_)
1419
+ sgr.rgb_dict = (
1420
+ self.ansi_format, {k: ~v for k, v in self._color_dict_.items()})
1421
+ return self._weak_var_update(_sgr_=sgr)
1422
+
1423
+ def __new__(cls, obj=None, color_spec=None, **kwargs):
1424
+ if ansi_type := kwargs.get('ansi_type'):
1425
+ ansi_type = get_ansi_type(ansi_type)
1426
+ if type(color_spec) is cls:
1427
+ if (ansi_type is not None
1428
+ and any(
1429
+ type(a) is not ansi_type
1430
+ for a in color_spec.ansi)):
1431
+ return color_spec.as_ansi_type(ansi_type)
1432
+ inst = super().__new__(cls, str(color_spec))
1433
+ for name, value in vars(color_spec).items():
1434
+ setattr(inst, name, value)
1435
+ return inst
1436
+ d = {'_ansi_type_': ansi_type or DEFAULT_ANSI}
1437
+ no_reset = d['_no_reset_'] = bool(kwargs.get('no_reset', False))
1438
+ suffix = '' if no_reset else SGR_RESET
1439
+ if obj is not None:
1440
+ if not isinstance(obj, str):
1441
+ obj = str(obj, encoding='ansi') if isinstance(obj, Buffer) else str(obj)
1442
+ if color_spec is None and obj.startswith(CSI.decode()):
1443
+ color_spec = obj.encode()
1444
+ obj = None
1445
+ elif color_spec is None:
1446
+ inst = super().__new__(cls, suffix)
1447
+ inst.__dict__ |= {'_sgr_': SgrSequence(), '_base_str_': str(), **d}
1448
+ return inst
1449
+ sgr, base_str_ = d['_sgr_'], d['_base_str_'] = (
1450
+ _get_color_str_vars(
1451
+ obj, color_spec, cast(AnsiColorType, ansi_type)))
1452
+ if ansi_type is None and sgr.is_color():
1453
+ d['_ansi_type_'], _ = max(
1454
+ Counter(type(p._value_) for p in sgr._sgr_params_ if p.is_color()).items(),
1455
+ key=op.itemgetter(1))
1456
+ inst = super().__new__(cls, ''.join([str(sgr), base_str_, suffix]))
1457
+ inst.__dict__ |= d
1458
+ return inst
1459
+
1460
+ def __repr__(self):
1461
+ return (f"{type(self).__name__}(%r, ansi_type=%s)"
1462
+ % (self.ansi.decode() + self.base_str,
1463
+ getattr(self.ansi_format, '__name__', type(None).__name__)))
1464
+
1465
+ def __sub__(self, other):
1466
+ """Return a copy of `self` with colors adjusted by color difference with `other`"""
1467
+ if (vt := type(other)) not in {Color, ColorStr}:
1468
+ raise TypeError(
1469
+ 'unsupported operand type(s) for -: '
1470
+ f"{ColorStr.__name__!r} and {vt.__qualname__!r}")
1471
+
1472
+ def _rgb_diff_color(a: Int3Tuple, b: Int3Tuple) -> Color:
1473
+ return Color.from_rgb(rgb_diff(a, b))
1474
+
1475
+ k: Literal['fg', 'bg']
1476
+ if vt is Color:
1477
+ diff_dict = {
1478
+ k: _rgb_diff_color(v, other.rgb)
1479
+ for k, v in self.rgb_dict.items()
1480
+ }
1481
+ else:
1482
+ diff_dict = {
1483
+ k: _rgb_diff_color(self.rgb_dict[k], other.rgb_dict[k])
1484
+ for k in self.rgb_dict.keys() & other.rgb_dict
1485
+ }
1486
+ if not diff_dict:
1487
+ return self
1488
+ sgr = SgrSequence(self._sgr_)
1489
+ sgr.rgb_dict = self.ansi_format, diff_dict
1490
+ return self._weak_var_update(_sgr_=sgr)
1491
+
1492
+ _ansi_ = _AnsiBytesGetter()
1493
+ _color_dict_ = _ColorDictGetter()
1494
+ _sgr_params_ = _SgrParamsGetter()
1495
+
1496
+ @property
1497
+ def ansi(self):
1498
+ return self._ansi_
1499
+
1500
+ @property
1501
+ def ansi_format(self):
1502
+ return self._ansi_type_
1503
+
1504
+ @property
1505
+ def base_str(self):
1506
+ """The non-ANSI part of the string"""
1507
+ return self._base_str_
1508
+
1509
+ @property
1510
+ def bg(self):
1511
+ """Background color"""
1512
+ return self._color_dict_.get('bg')
1513
+
1514
+ @property
1515
+ def fg(self):
1516
+ """Foreground color"""
1517
+ return self._color_dict_.get('fg')
1518
+
1519
+ @property
1520
+ def no_reset(self):
1521
+ return self._no_reset_
1522
+
1523
+ @property
1524
+ def rgb_dict(self):
1525
+ return {k: v.rgb for k, v in self._color_dict_.items()}
1526
+
1527
+
1528
+ def hsl_gradient(start: Int3Tuple | Float3Tuple,
1529
+ stop: Int3Tuple | Float3Tuple,
1530
+ step: SupportsIndex,
1531
+ num: SupportsIndex = None,
1532
+ ncycles: int | float = float('inf'),
1533
+ replace_idx: tuple[
1534
+ SupportsIndex | Iterable[SupportsIndex],
1535
+ Iterator[Color]] = None,
1536
+ dtype: type[Color] | Callable[[Int3Tuple], int] = Color):
1537
+ replace_idx, rgb_iter = _resolve_replacement_indices(replace_idx)
1538
+ while abs(float(step)) < 1:
1539
+ step *= 10
1540
+ color_vec = _init_gradient_color_vec(num, start, step, stop)
1541
+ color_iter = iter(color_vec)
1542
+ type_map: dict[type[Color | int], ...] = {Color: lambda x: x.rgb, int: lambda x: hex2rgb(x)}
1543
+ get_rgb_iter_idx: Callable[[Color | int, SupportsIndex], int] = lambda x, ix: \
1544
+ rgb2hsl(type_map[type(x)](x))[ix]
1545
+ next_rgb_iter = None
1546
+ prev_output = None
1547
+ while ncycles > 0:
1548
+ try:
1549
+ cur_iter = next(color_iter)
1550
+ if cur_iter != prev_output:
1551
+ for idx in replace_idx:
1552
+ try:
1553
+ next_rgb_iter = next(rgb_iter)
1554
+ cur_iter = list(cur_iter)
1555
+ cur_iter[idx] = get_rgb_iter_idx(next_rgb_iter, idx)
1556
+ except StopIteration:
1557
+ raise GeneratorExit
1558
+ except KeyError:
1559
+ raise TypeError(
1560
+ f"Expected iterator to return "
1561
+ f"{repr(Color.__qualname__)} or {repr(int.__qualname__)}, "
1562
+ f"got {repr(type(next_rgb_iter).__qualname__)} instead") from None
1563
+ output = hsl2rgb(cast(Float3Tuple, cur_iter))
1564
+ if callable(dtype):
1565
+ output = dtype(output)
1566
+ yield output
1567
+ prev_output = cur_iter
1568
+ except StopIteration:
1569
+ ncycles -= 1
1570
+ color_vec.reverse()
1571
+ color_iter = iter(color_vec)
1572
+ except GeneratorExit:
1573
+ break
1574
+
1575
+
1576
+ def _resolve_replacement_indices(
1577
+ replace_idx: tuple[SupportsIndex | Sequence[SupportsIndex], Iterator[Color]] = None
1578
+ ):
1579
+ if replace_idx is not None:
1580
+ replace_idx, rgb_iter = replace_idx
1581
+ if not isinstance(rgb_iter, Iterator):
1582
+ raise TypeError(
1583
+ f"Expected 'replace_idx[1]' to be an iterator, got {type(rgb_iter).__name__} "
1584
+ f"instead")
1585
+ if not isinstance(replace_idx, Sequence):
1586
+ replace_idx = {replace_idx}
1587
+ else:
1588
+ replace_idx = set(replace_idx)
1589
+ valid_idx_range = range(3)
1590
+ if any(idx_diff := replace_idx.difference(valid_idx_range)):
1591
+ raise ValueError(
1592
+ f"Invalid replacement indices: {idx_diff}")
1593
+ if replace_idx == set(valid_idx_range):
1594
+ raise ValueError(
1595
+ f"All 3 indexes selected for replacement: {replace_idx=}")
1596
+ else:
1597
+ rgb_iter = None
1598
+ replace_idx = []
1599
+ return replace_idx, rgb_iter
1600
+
1601
+
1602
+ def _init_gradient_color_vec(num: SupportsIndex,
1603
+ start: Int3Tuple | Float3Tuple,
1604
+ step: SupportsIndex,
1605
+ stop: Int3Tuple | Float3Tuple):
1606
+ def convert_bounds(rgb: Int3Tuple):
1607
+ if all(0 <= n <= 255 for n in rgb):
1608
+ return rgb2hsl(rgb)
1609
+ raise ValueError
1610
+
1611
+ start, stop = tuple(map(convert_bounds, (start, stop)))
1612
+ start_h, start_s, start_l = start
1613
+ stop_h, stop_s, stop_l = stop
1614
+ if num:
1615
+ num_samples = num
1616
+ else:
1617
+ abs_h = abs(stop_h - start_h)
1618
+ h_diff = min(abs_h, 360 - abs_h)
1619
+ dist = math.sqrt(h_diff ** 2 + (stop_s - start_s) ** 2 + (stop_l - start_l) ** 2)
1620
+ num_samples = max(int(dist / float(step)), 1)
1621
+ color_vec = [np.linspace(*bounds, num=num_samples, dtype=float) for bounds in zip(start, stop)]
1622
+ color_vec = list(zip(*color_vec))
1623
+ return color_vec
1624
+
1625
+
1626
+ def rgb_luma_transform(rgb: Int3Tuple,
1627
+ start: SupportsIndex = None,
1628
+ num: SupportsIndex = 50,
1629
+ step: SupportsIndex = 1,
1630
+ cycle: bool | Literal['wave'] = False,
1631
+ ncycles: int | float = float('inf'),
1632
+ gradient: Int3Tuple = None,
1633
+ dtype: type[Color] = None) -> Iterator[Int3Tuple | int | Color]:
1634
+ if dtype is None:
1635
+ ret_type = tuple
1636
+ elif issubclass(dtype, int):
1637
+ ret_type = lambda x: dtype(rgb2hex(x))
1638
+ is_cycle = bool(cycle is not False)
1639
+ is_oscillator = cycle == 'wave'
1640
+ if is_oscillator:
1641
+ ncycles *= 2
1642
+ h, s, luma = rgb2hsl(rgb)
1643
+ luma_linspace = [*np.linspace(start=0, stop=1, num=num)][::step]
1644
+ if start:
1645
+ start = min(max(float(start), 0), 1)
1646
+ luma = min(luma_linspace, key=lambda x: abs(x - start))
1647
+ start_idx = luma_linspace.index(luma)
1648
+ remaining_indices = luma_linspace[start_idx:]
1649
+ luma_iter = iter(remaining_indices)
1650
+ else:
1651
+ luma_iter = iter(luma_linspace)
1652
+
1653
+ def _generator():
1654
+ nonlocal luma_iter, ncycles
1655
+ if step == 0:
1656
+ yield rgb
1657
+ return
1658
+ prev_output = None
1659
+ while ncycles > 0:
1660
+ try:
1661
+ output = hsl2rgb((h, s, next(luma_iter)))
1662
+ if output != prev_output:
1663
+ yield ret_type(output)
1664
+ prev_output = output
1665
+ except StopIteration as STOP_IT:
1666
+ if not is_cycle:
1667
+ raise STOP_IT
1668
+ ncycles -= 1
1669
+ if is_oscillator:
1670
+ luma_linspace.reverse()
1671
+ luma_iter = iter(luma_linspace)
1672
+
1673
+ if gradient is not None:
1674
+ _gradient = hsl_gradient(
1675
+ start=rgb, stop=gradient, step=step, num=num, replace_idx=(2, _generator()))
1676
+ return iter(_gradient)
1677
+ return iter(_generator())