File: lpy.py
   1 #!/usr/bin/python
   2 
   3 # The MIT License (MIT)
   4 #
   5 # Copyright (c) 2026 pacman64
   6 #
   7 # Permission is hereby granted, free of charge, to any person obtaining a copy
   8 # of this software and associated documentation files (the "Software"), to deal
   9 # in the Software without restriction, including without limitation the rights
  10 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  11 # copies of the Software, and to permit persons to whom the Software is
  12 # furnished to do so, subject to the following conditions:
  13 #
  14 # The above copyright notice and this permission notice shall be included in
  15 # all copies or substantial portions of the Software.
  16 #
  17 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  18 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  19 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  20 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  21 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  22 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  23 # SOFTWARE.
  24 
  25 
  26 info = '''
  27 lpy [options...] [python expression] [files/URIs...]
  28 
  29 
  30 Legend using PYthon colors lines by running the Python expression given on
  31 each line, and using its common results. Inputs are encoded as UTF-8. Any
  32 ANSI-styles are ignored from the input, to ensure consistent output-styling.
  33 
  34 Carriage-returns are always ignored in lines, as well as any UTF-8-BOM on the
  35 first line of each input.
  36 
  37 The expression can use either `l` or `line` for the current line, and `i` as
  38 a 0-based line counter which keeps growing even across input-sources, when
  39 given multiple inputs. Also available is `n`, a 1-based line counter which
  40 otherwise works the same way.
  41 
  42 Each line is automatically parsed as JSON: when successful, the parsed line
  43 is available to the expression as `v`, or `value`. You can check failure to
  44 parse JSON by checking if `v` or `value` is of type Skip, since None can be
  45 the result of successfully parsing a null JSON value.
  46 
  47 Input-sources can be either files or web-URIs. When not given any explicit
  48 named sources, the standard input is used. A single dash also means the
  49 standard input.
  50 
  51 When the expression results in None, the current input line is ignored,
  52 allowing expressions to also act as convenient line filters.
  53 
  54 
  55 Options
  56 
  57 All options can start with either a single or a double leading dash:
  58 
  59     -h, -help                      show this help message
  60     -m, -mod, -module, -modules    import modules named in the next argument,
  61                                    where multiple names are comma-separated
  62     -t, -trace, -traceback         turn exceptions into multi-line tracebacks
  63     -v, -verbose                   also show legend-keys to stderr at the end
  64 
  65 
  66 Examples
  67 
  68 # read from a file, and only emit (equally-colored) lines exceeding 80 items
  69 lpy '1 if len(line) > 80 else None' data.txt
  70 
  71 # color files in current folder (roughly) by common auto-detected file-type
  72 file * | lpy '(lambda e: e[1] + e[2])(" ".join(l.split(":")[1:]).split())'
  73 
  74 # color files in current folder (roughly) by common auto-detected file-type
  75 file * | lpy 'fields[1] if len(fields) < 3 else fields[1] + fields[2]'
  76 
  77 # color numbers by common remainder when divided by 5
  78 seq 50 | lpy 'v % 5'
  79 
  80 # transform lines in order with a format-string, styling nothing
  81 seq 50 | lpy 'print(f"{line + "**2":>6} = {float(v**2):12,.4f}")'
  82 
  83 # color numbers by common remainder when divided by 5, ignoring numbers > 30
  84 seq 50 | lpy 'None if v > 30 else v % 5'
  85 '''
  86 
  87 
  88 from itertools import islice, zip_longest
  89 from json import dump, dumps, loads
  90 from math import isinf, isnan
  91 from os import getenv
  92 from re import compile as compile_uncached, IGNORECASE
  93 from sys import argv, exit, stderr, stdin, stdout
  94 from time import localtime, strftime
  95 from typing import Iterable
  96 
  97 
  98 if len(argv) > 1 and argv[1] in ('-h', '--h', '-help', '--help'):
  99     print(info.strip())
 100     exit(0)
 101 
 102 
 103 def style_lines(src, expr):
 104     # `comprehension` expressions seem to ignore local variables: even
 105     # lambda-based workaround-attempts fail to make needed values like
 106     # the current line available to such expressions
 107     global i, n, fn, l, line, v, val, value
 108     global legend, styles, resets, items, fields, words, tsv
 109 
 110     fn = 1
 111 
 112     for l in src:
 113         l = l.rstrip('\r\n').rstrip('\n')
 114         if fn == 1:
 115             l = l.lstrip('\xef\xbb\xbf')
 116 
 117         line = l
 118         items = fields = words = line.split()
 119         tsv = line.split('\t')
 120 
 121         try:
 122             v = val = value = loads(l)
 123         except Exception as _:
 124             v = val = value = Skip()
 125 
 126         k = eval(expr)
 127         if callable(k) and (not isinstance(k, Skip)):
 128             k = k(line)
 129         if isinstance(k, BaseException):
 130             raise k
 131 
 132         i += 1
 133         n += 1
 134         fn += 1
 135 
 136         if (k is None) or isinstance(k, Skip):
 137             continue
 138         if not (k in legend):
 139             legend[k] = len(legend)
 140 
 141         k = legend[k] % len(styles)
 142         # out_line = line.replace('\x1b[0m', resets[k])
 143         out_line = plain(line)
 144         print(f'{styles[k]}{out_line}\x1b[0m', flush=flushed)
 145 
 146 def fail(msg, code = 1):
 147     print(str(msg), file=stderr)
 148     exit(code)
 149 
 150 def make_open_utf8(open):
 151     def open_utf8_readonly(path):
 152         return open(path, encoding='utf-8')
 153     return open_utf8_readonly
 154 
 155 def seemsurl(path):
 156     protocols = ('https://', 'http://', 'file://', 'ftp://', 'data:')
 157     return any(path.startswith(p) for p in protocols)
 158 
 159 class Skip:
 160     def __call__(self, x):
 161         return isinstance(x, self.__class__)
 162 
 163 skip = Skip()
 164 
 165 def cases(x, *args):
 166     '''
 167     Simulate a switch statement on a value, using matches/result pairs from
 168     the arguments given; when given an even number of extra args, None is
 169     used as a final fallback result; when given an odd number of extra args,
 170     the last argument is used as a final `default` value, if needed.
 171     '''
 172 
 173     for i in range(0, len(args) - len(args) % 2, 2):
 174         test, res = args[i], args[i+1]
 175         if isinstance(test, (list, tuple)) and (x in test):
 176             return res
 177         if isinstance(test, float) and isnan(test) and isnan(x):
 178             return res
 179         if x == test:
 180             return res
 181     return None if len(args) % 2 == 0 else args[-1]
 182 
 183 switch = cases
 184 
 185 def chunk(items, chunk_size):
 186     'Break iterable into chunks, each with up to the item-count given.'
 187 
 188     if isinstance(items, str):
 189         n = len(items)
 190         while n >= chunk_size:
 191             yield items[:chunk_size]
 192             items = items[chunk_size:]
 193             n -= chunk_size
 194         if n > 0:
 195             yield items
 196         return
 197 
 198     if not isinstance(chunk_size, int):
 199         raise Exception('non-integer chunk-size')
 200     if chunk_size < 1:
 201         raise Exception('non-positive chunk-size')
 202 
 203     it = iter(items)
 204     while True:
 205         head = tuple(islice(it, chunk_size))
 206         if not head:
 207             return
 208         yield head
 209 
 210 chunked = chunk
 211 
 212 # re_cache is used by custom func compile to cache previously-compiled
 213 # regular-expressions, which makes them quicker to (re)use in formulas
 214 re_cache = {}
 215 
 216 def re_compile(expr, flags = 0):
 217     'Speed-up using regexes across lines, by avoiding recompilations.'
 218 
 219     if flags in re_cache:
 220         cache = re_cache[flags]
 221     else:
 222         cache = {}
 223         re_cache[flags] = cache
 224     if expr in cache:
 225         return cache[expr]
 226 
 227     pat = compile_uncached(expr, flags)
 228     cache[expr] = pat
 229     return pat
 230 
 231 def icompile(expr):
 232     return re_compile(expr, IGNORECASE)
 233 
 234 def cond(*args):
 235     for i, e in enumerate(args):
 236         if i % 2 == 0 and i < len(args) - 1 and e:
 237             return args[i + 1]
 238 
 239     return args[-1] if len(args) % 2 == 1 else None
 240 
 241 def denan(x, fallback = None):
 242     'Replace floating-point NaN with the alternative value given.'
 243     return x if not (isinstance(x, float) and isnan(x)) else fallback
 244 
 245 def denil(*args):
 246     'Avoid None values, if possible: first value which isn\'t None wins.'
 247     for e in args:
 248         if e != None:
 249             return e
 250     return None
 251 
 252 denone = denull = denil
 253 
 254 def dive(into, using):
 255     'Depth-first recursive caller for 1-input functions.'
 256 
 257     if callable(into):
 258         into, using = using, into
 259 
 260     def rec(v):
 261         if isinstance(v, dict):
 262             return {k: rec(v) for k, v in v.items()}
 263         if isinstance(v, Iterable) and not isinstance(v, str):
 264             return [rec(v) for v in v]
 265         return using(v)
 266 
 267     return rec(into)
 268 
 269 def divekeys(into, using):
 270     'Depth-first recursive caller for 2-input funcs which rename dict keys.'
 271 
 272     if callable(into):
 273         into, using = using, into
 274 
 275     def rec(v):
 276         if isinstance(v, dict):
 277             return {using(k): rec(v) for k, v in v.items()}
 278         if isinstance(v, Iterable) and not isinstance(v, str):
 279             return [rec(v) for i, v in enumerate(v)]
 280         return v
 281 
 282     return rec(None, into)
 283 
 284 def divekv(into, using, using2 = None):
 285     'Depth-first recursive caller for 2-input functions.'
 286 
 287     if using2 is None:
 288         if callable(into):
 289             into, using = using, into
 290     else:
 291         if not callable(using2):
 292             into, using, using2 = using2, into, using
 293 
 294     def rec(k, v):
 295         if isinstance(v, dict):
 296             return {k: rec(k, v) for k, v in v.items()}
 297         if isinstance(v, Iterable) and not isinstance(v, str):
 298             return [rec(i, v) for i, v in enumerate(v)]
 299         return using(k, v)
 300 
 301     def rec2(k, v):
 302         if isinstance(v, dict):
 303             return {str(using(k, v)): rec2(k, v) for k, v in v.items()}
 304         if isinstance(v, Iterable) and not isinstance(v, str):
 305             return [rec2(i, v) for i, v in enumerate(v)]
 306         return using2(k, v)
 307 
 308     return rec(None, into) if using2 is None else rec2(None, into)
 309 
 310 kvdive = divekv
 311 
 312 def drop(src, *what):
 313     if isinstance(src, str):
 314         for s in what:
 315             src = src.replace(s, '')
 316         return src
 317 
 318     def kdrop(src, what):
 319         return {k: v for (k, v) in src.items() if not (k in what)}
 320 
 321     if isinstance(src, dict):
 322         return kdrop(src, set(what))
 323 
 324     if isinstance(src, Iterable):
 325         what = set(what)
 326         return [kdrop(e, what) for e in src if isinstance(e, dict)]
 327 
 328     return None
 329 
 330 dropped = drop
 331 
 332 def flat(*args):
 333     'Turn arbitrarily-nested values/sequences into a single flat sequence.'
 334 
 335     def rec(x):
 336         if x is None:
 337             return
 338 
 339         if isinstance(x, dict):
 340             yield from rec(x.values())
 341 
 342         if isinstance(x, str):
 343             yield x
 344             return
 345 
 346         if isinstance(x, Iterable):
 347             for e in x:
 348                 yield from rec(e)
 349             return
 350 
 351         yield x
 352 
 353     for x in args:
 354         yield from rec(x)
 355 
 356 flatten = flattened = flat
 357 
 358 def identity(x):
 359     return x
 360 
 361 def iota(n):
 362     return range(1, n + 1)
 363 
 364 def join(x, *y):
 365     'Join values into a string, or make a dict from keys and values.'
 366 
 367     if len(y) == 0:
 368         return ' '.join(str(v) for v in x)
 369     if isinstance(x, str):
 370         return x.join(str(v) for v in y)
 371     if len(y) == 1 and isinstance(y[0], str):
 372         return y[0].join(str(v) for v in x)
 373     if len(y) == 1 and isinstance(y[0], (list, range, set, tuple, Generator)):
 374         return {k: v for k, v in zip_longest(x, y[0]) if not (k is None)}
 375     if isinstance(y, (list, range, set, tuple, Generator)):
 376         return {k: v for k, v in zip_longest(x, y) if not (k is None)}
 377     return {k: y for k in x}
 378 
 379 joined = join
 380 
 381 def maybe(f, x = None):
 382     try:
 383         return f(x)
 384     except Exception as _:
 385         return x
 386 
 387 def number(x):
 388     try:
 389         return int(x)
 390     except Exception as _:
 391         pass
 392     try:
 393         return float(x)
 394     except Exception as _:
 395         return x
 396 
 397 def pick(src, *keys):
 398     if isinstance(src, dict):
 399         return {k: src.get(k, None) for k in keys}
 400     return [{k: e.get(k, None) for k in keys} for e in src if isinstance(e, dict)]
 401 
 402 picked = pick
 403 
 404 def plain(s):
 405     'Ignore all ANSI-style sequences in a string.'
 406     return re_compile('''\x1b\\[([0-9;]+m|[0-9]*[A-HJKST])''').sub('', s)
 407 
 408 def predicate(x):
 409     'Helps various higher-order funcs, by standardizing `predicate` values.'
 410     if callable(x):
 411         return x
 412     if not isinstance(x, float):
 413         return lambda y: x == y
 414     if isnan(x):
 415         return lambda y: isinstance(y, float) and isnan(y)
 416     if isinf(x):
 417         return lambda y: isinstance(y, float) and isinf(y)
 418 
 419 def rescue(attempt, fallback = None):
 420     try:
 421         return attempt()
 422     except Exception as e:
 423         if callable(fallback) and (not isinstance(fallback, Skip)):
 424             return fallback(e)
 425         return fallback
 426 
 427 rescued = rescue
 428 
 429 def retype(x):
 430     'Try to narrow the type of the value given.'
 431 
 432     if isinstance(x, float):
 433         n = int(x)
 434         return n if float(n) == x else x
 435 
 436     if not isinstance(x, str):
 437         return x
 438 
 439     try:
 440         return loads(x)
 441     except Exception as _:
 442         pass
 443 
 444     try:
 445         return int(x)
 446     except Exception as _:
 447         pass
 448 
 449     try:
 450         return float(x)
 451     except Exception as _:
 452         pass
 453 
 454     return x
 455 
 456 autocast = autocasted = mold = molded = recast = recasted = remold = retype
 457 remolded = retyped = retype
 458 
 459 def json0(x):
 460     if isinstance(x, (range, set, Generator)):
 461         x = tuple(x)
 462     return dumps(x, separators=(',', ':'), allow_nan=False, indent=None)
 463 
 464 j0 = json0
 465 
 466 def typeof(x):
 467     return {
 468         type(None): 'null',
 469         bool: 'boolean',
 470         dict: 'object',
 471         float: 'number', int: 'number',
 472         str: 'string',
 473         list: 'array', tuple: 'array',
 474     }.get(type(x), 'other')
 475 
 476 jstype = typeof
 477 
 478 def after(x, what):
 479     i = x.find(what)
 480     return '' if i < 0 else x[i+len(what):]
 481 
 482 def afterlast(x, what):
 483     i = x.rfind(what)
 484     return '' if i < 0 else x[i+len(what):]
 485 
 486 afterfinal = afterlast
 487 
 488 def before(x, what):
 489     i = x.find(what)
 490     return x if i < 0 else x[:i]
 491 
 492 def beforelast(x, what):
 493     i = x.rfind(what)
 494     return x if i < 0 else x[:i]
 495 
 496 beforefinal = beforelast
 497 
 498 def since(x, what):
 499     i = x.find(what)
 500     return '' if i < 0 else x[i:]
 501 
 502 def sincelast(x, what):
 503     i = x.rfind(what)
 504     return '' if i < 0 else x[i:]
 505 
 506 sincefinal = sincelast
 507 
 508 def until(x, what):
 509     i = x.find(what)
 510     return x if i < 0 else x[:i+len(what)]
 511 
 512 def untilfinal(x, what):
 513     i = x.rfind(what)
 514     return x if i < 0 else x[:i+len(what)]
 515 
 516 untillast = untilfinal
 517 
 518 def message(msg, result = None):
 519     print(msg, file=stderr)
 520     return result
 521 
 522 msg = message
 523 
 524 # seen is used by func `once` to remember previously-given values
 525 seen = set()
 526 
 527 def once(x, key = None):
 528     k = key(x) if key else x
 529     if k in seen:
 530         return None
 531     seen.add(k)
 532     return x
 533 
 534 dedup = unique = once
 535 
 536 def utf8(x):
 537     try:
 538         if isinstance(x, str):
 539             x = x.encode('utf-8')
 540         return str(x, 'utf-8')
 541     except Exception as _:
 542         return None
 543 
 544 def ymdhms(when = None):
 545     fmt = f'%Y-%m-%d %H:%M:%S'
 546     if isinstance(when, (float, int)):
 547         return strftime(fmt, localtime(float(when)))
 548     if isinstance(when, tuple):
 549         return strftime(fmt, when)
 550     return strftime(fmt, localtime())
 551 
 552 
 553 amp = '&'
 554 apos = '\''
 555 lcur = lcurly = '{'
 556 quot = '"'
 557 rcur = rcurly = '}'
 558 
 559 true = True
 560 false = False
 561 nil = none = null = None
 562 
 563 
 564 exec = None
 565 open_utf8 = make_open_utf8(open)
 566 open = open_utf8
 567 
 568 modules_opts = (
 569     '-m', '--m', '-mod', '--mod', '-module', '--module',
 570     '-modules', '--modules',
 571 )
 572 trace_opts = ('-t', '--t', '-trace', '--trace', '-traceback', '--traceback')
 573 verbose_opts = ('-v', '--v', '-verbose', '--verbose')
 574 
 575 args = argv[1:]
 576 if any(seemsurl(e) for e in args):
 577     from io import TextIOWrapper
 578     from urllib.request import urlopen
 579 
 580 tracing = False
 581 verbose = False
 582 
 583 while len(args) > 0:
 584     if args[0] == '-':
 585         break
 586 
 587     if args[0] == '--':
 588         args = args[1:]
 589         break
 590 
 591     if args[0] in modules_opts:
 592         try:
 593             if len(args) < 2:
 594                 msg = 'a module name or a comma-separated list of modules'
 595                 raise Exception('expected ' + msg)
 596 
 597             g = globals()
 598             from importlib import import_module
 599             for e in args[1].split(','):
 600                 g[e] = import_module(e)
 601 
 602             g = None
 603             import_module = None
 604             args = args[2:]
 605         except Exception as e:
 606             fail(e, 1)
 607 
 608         continue
 609 
 610     if args[0] in verbose_opts:
 611         verbose = True
 612         args = args[1:]
 613         continue
 614 
 615     if args[0] in trace_opts:
 616         tracing = True
 617         args = args[1:]
 618         continue
 619 
 620     if args[0].startswith('-'):
 621         fail(f'lpy: unsupported option {args[0]}', 1)
 622         exit(1)
 623 
 624     break
 625 
 626 
 627 # ensure live-lines output, unless stdout is being saved into a file
 628 flushed = stdout.isatty() or (not stdout.seekable())
 629 
 630 # '\x1b[38;2;255;135;255m', # pink (originally after cyan)
 631 
 632 styles = (
 633     '\x1b[38;2;0;95;215m', # blue
 634     '\x1b[38;2;215;95;0m', # orange
 635     '\x1b[38;2;135;95;255m', # purple
 636     '\x1b[38;2;0;175;215m', # cyan
 637     '\x1b[38;2;255;85;255m', # pink
 638     '\x1b[38;2;0;135;95m', # green
 639     '\x1b[38;2;204;0;0m', # red
 640     '\x1b[38;2;168;168;168m', # gray
 641 )
 642 
 643 if (getenv('COLORBLIND', '0') != '0') or (getenv('COLOR_BLIND', '0') != '0'):
 644     styles = (
 645         '\x1b[38;2;0;95;215m', # blue
 646         '\x1b[38;2;215;95;0m', # orange
 647         '\x1b[38;2;135;95;255m', # purple
 648         '\x1b[38;2;0;175;215m', # cyan
 649         '\x1b[38;2;255;85;255m', # pink
 650         '\x1b[38;2;0;135;95m', # green
 651         '\x1b[38;2;168;168;168m', # gray
 652     )
 653 getenv = None
 654 
 655 resets = tuple('\x1b[0m' + s for s in styles)
 656 
 657 expr = '.'
 658 if len(args) > 0:
 659     expr = args[0]
 660     args = args[1:]
 661 
 662 if expr == '.':
 663     expr = 'line'
 664 
 665 try:
 666     expr = compile(expr, expr, mode='eval')
 667     compile = None
 668 
 669     i = 0
 670     n = 1
 671     v = val = value = Skip()
 672 
 673     legend = {}
 674     items = fields = words = []
 675     tsv = []
 676 
 677     if len(args) == 0:
 678         style_lines(stdin, expr)
 679 
 680     if args.count('-') > 1:
 681         msg = 'reading from `-` (standard input) more than once not allowed'
 682         raise ValueError(msg)
 683 
 684     for path in args:
 685         if path == '-':
 686             style_lines(stdin, expr)
 687             continue
 688 
 689         if seemsurl(path):
 690             with urlopen(path) as inp:
 691                 with TextIOWrapper(inp, encoding='utf-8') as txt:
 692                     style_lines(txt, expr)
 693             continue
 694 
 695         with open_utf8(path) as txt:
 696             style_lines(txt, expr)
 697 
 698     if verbose:
 699         for k, v in legend.items():
 700             print(f'{styles[v]}{k}\x1b[0m', file=stderr)
 701 except BrokenPipeError:
 702     # quit quietly, instead of showing a confusing error message
 703     stderr.close()
 704     exit(0)
 705 except KeyboardInterrupt:
 706     exit(2)
 707 except Exception as e:
 708     if tracing:
 709         raise e
 710     fail(e, 1)