File: zj.py
   1 #!/usr/bin/python3
   2 
   3 # The MIT License (MIT)
   4 #
   5 # Copyright © 2024 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 from inspect import getfullargspec
  27 from itertools import islice
  28 from json import load, dump
  29 from sys import argv, stderr, stdin, stdout
  30 from typing import Any, Callable, Dict, Iterable, NoReturn, Tuple
  31 
  32 
  33 # extra imports for the `python-lambda` option
  34 
  35 from decimal import Decimal, getcontext
  36 
  37 from fractions import Fraction
  38 
  39 import functools
  40 from functools import \
  41     cache, cached_property, cmp_to_key, get_cache_token, lru_cache, \
  42     namedtuple, partial, partialmethod, recursive_repr, reduce, \
  43     singledispatch, singledispatchmethod, total_ordering, update_wrapper, \
  44     wraps
  45 
  46 import itertools
  47 from itertools import \
  48     accumulate, chain, combinations, combinations_with_replacement, \
  49     compress, count, cycle, dropwhile, filterfalse, groupby, islice, \
  50     permutations, product, repeat, starmap, takewhile, tee, zip_longest
  51 try:
  52     from itertools import pairwise
  53     from itertools import batched
  54 except Exception:
  55     pass
  56 
  57 from json import dumps, loads
  58 
  59 import math
  60 Math = math
  61 from math import \
  62     acos, acosh, asin, asinh, atan, atan2, atanh, ceil, comb, \
  63     copysign, cos, cosh, degrees, dist, e, erf, erfc, exp, expm1, \
  64     fabs, factorial, floor, fmod, frexp, fsum, gamma, gcd, hypot, inf, \
  65     isclose, isfinite, isinf, isnan, isqrt, lcm, ldexp, lgamma, log, \
  66     log10, log1p, log2, modf, nan, nextafter, perm, pi, pow, prod, \
  67     radians, remainder, sin, sinh, sqrt, tan, tanh, tau, trunc, ulp
  68 try:
  69     from math import cbrt, exp2
  70 except Exception:
  71     pass
  72 
  73 power = pow
  74 
  75 import operator
  76 
  77 import statistics
  78 from statistics import \
  79     bisect_left, bisect_right, fmean, \
  80     geometric_mean, harmonic_mean, mean, median, \
  81     median_grouped, median_high, median_low, mode, multimode, pstdev, \
  82     pvariance, quantiles, stdev, variance
  83 try:
  84     from statistics import \
  85         correlation, covariance, linear_regression, mul
  86 except Exception:
  87     pass
  88 
  89 import string
  90 from string import \
  91     Formatter, Template, ascii_letters, ascii_lowercase, ascii_uppercase, \
  92     capwords, digits, hexdigits, octdigits, printable, punctuation, \
  93     whitespace
  94 
  95 alphabet = ascii_letters
  96 letters = ascii_letters
  97 lowercase = ascii_lowercase
  98 uppercase = ascii_uppercase
  99 
 100 from textwrap import dedent, fill, indent, shorten, wrap
 101 
 102 from urllib.parse import \
 103     parse_qs, parse_qsl, quote, quote_from_bytes, quote_plus, unquote, \
 104     unquote_plus, unquote_to_bytes, unwrap, urldefrag, urlencode, urljoin, \
 105     urlparse, urlsplit, urlunparse, urlunsplit
 106 
 107 
 108 from re import compile as zj_compile_re
 109 
 110 
 111 zj_info_msg = '''
 112 zj [keys/indices...]
 113 
 114 
 115 Zoom Json digs into a subset of valid JSON input, using the given mix of
 116 keys and array-indices, the latter being either 0-based or negative, to
 117 index backward from the ends of arrays.
 118 
 119 Zooming on object keys is first tried as an exact key-match, failing that
 120 as a case-insensitive key-match (first such match): when both approaches
 121 fail, if the key is a valid integer, the key at the (even negative) index
 122 given is used.
 123 
 124 Invalid array-indices and missing object-keys result in null values, when
 125 none of the special keys/fallbacks shown later apply.
 126 
 127 You can slice arrays the exclusive/go/python way using index-pairs with a
 128 `:` between the start/end pair, as long as it's a single argument; you can
 129 even use `..` as the index-pair separator to include the stop index in the
 130 result. Either way, as with go/python, you can omit either of the indices
 131 when slicing.
 132 
 133 Special key `.` acts as implicit loops on arrays, and even objects without
 134 that specific key: in the unlikely case that an object has `.` as one of
 135 its keys, you can use one of loop-fallback aliases, shown later.
 136 
 137 Another special key is `+` (no quotes): when used, the rest of the keys
 138 are used `in parallel`, allowing multiple picks from the current value.
 139 When picking array items, you can also use either type (`:` or `..`) of
 140 slicing, even mixing it with individual indices.
 141 
 142 Similar to `+`, the `-` fallback-key drops keys, which means all items are
 143 picked, except for those mentioned after the `-`.
 144 
 145 Unlike the looping special key, after the first `+` special-key, all keys
 146 following it, special or not, are picked normally.
 147 
 148 In case any of the special keys are actual keys in the data loaded, some
 149 aliases are available:
 150 
 151     .   /.  ./  :.  .:
 152     +   /+  +/  :+  +:
 153     -   /-  -/  :-  -:
 154 
 155     .i   :i   .info    :info     :info:
 156     .k   :k   .keys    :keys     :keys:
 157     .t   :t   .type    :type     :type:
 158     .l   :l   .len     .length   :len      :len:    :length    :length:
 159 
 160 These aliases allow using the special functionality even on objects whose
 161 keys match some of these special names, as it's extremely unlikely data use
 162 all aliases as actual keys at any level.
 163 
 164 The only input supported is valid JSON coming from standard-input: there's
 165 no way to load files using their names. To load data from files/URIs use
 166 tools like `cat` or `curl`, and pipe their output into this tool.
 167 '''
 168 
 169 zj_slice_re = zj_compile_re('''^(([+-]?[0-9]+)?)(:|\.\.)(([+-]?[0-9]+)?)$''')
 170 
 171 
 172 def zj_zoom(data: Any, keys: Tuple[str, ...]) -> Any:
 173     eval_due = False
 174 
 175     for i, k in enumerate(keys):
 176         try:
 177             if eval_due:
 178                 data = eval(k)(data)
 179                 eval_due = False
 180                 continue
 181 
 182             if isinstance(data, dict):
 183                 m = zj_match_key(data, k)
 184                 if m in data:
 185                     data = data[m]
 186                     continue
 187                 m = zj_slice_re.match(k)
 188                 if m:
 189                     data = {k: data[k] for k in zj_match_keys(data, k)}
 190                     continue
 191 
 192             if isinstance(data, (list, tuple)):
 193                 try:
 194                     i = int(k)
 195                     l = len(data)
 196                     data = data[i] if -l <= i < l else None
 197                     continue
 198                 except Exception:
 199                     m = zj_slice_re.match(k)
 200                     if m:
 201                         data = [data[i] for i in zj_match_indices(data, k)]
 202                         continue
 203 
 204             if k in ('.pyl', ':pyl', 'pyl:', ':pyl:'):
 205                 eval_due = True
 206                 continue
 207 
 208             if k in ('.', '/.', './', ':.', '.:'):
 209                 if isinstance(data, dict):
 210                     rest = tuple(keys[i + 1:])
 211                     return {k: zj_zoom(v, rest) for k, v in data.items()}
 212                 if isinstance(data, (list, tuple)):
 213                     rest = tuple(keys[i + 1:])
 214                     return tuple(zj_zoom(v, rest) for v in data)
 215 
 216                 # doing nothing amounts to an identity-op for simple values
 217                 continue
 218 
 219             fn = zj_final_fallbacks.get(k, None)
 220             if fn:
 221                 return fn(data, tuple(keys[i + 1:]))
 222 
 223             fn = zj_fallbacks.get(k, None)
 224             if fn:
 225                 data = fn(data)
 226                 continue
 227 
 228             if isinstance(data, (dict, list, tuple)):
 229                 data = None
 230                 continue
 231 
 232             kind = zj_type(data)
 233             msg = f'value of type {kind} has no properties to zoom into'
 234             raise Exception(msg)
 235         except Exception as e:
 236             key_path = ' > '.join(islice(keys, None, i + 1))
 237             raise Exception(f'{key_path}: {e}')
 238 
 239     return data
 240 
 241 
 242 def zj_match_key(src: Dict, key: str) -> str:
 243     if key in src:
 244         return key
 245 
 246     low = key.casefold()
 247     for k in src.keys():
 248         if low == k.casefold():
 249             return k
 250 
 251     try:
 252         i = int(key)
 253         l = len(src)
 254         if i < 0:
 255             i += l
 256         if i < 0 or i >= l:
 257             return None
 258         for j, k in enumerate(src.keys()):
 259             if i == j:
 260                 return k
 261     except Exception:
 262         return key
 263     return key
 264 
 265 
 266 def zj_match_keys(src: Any, key: str) -> Iterable:
 267     if isinstance(src, (list, tuple)):
 268         yield from zj_match_indices(src, key)
 269         yield from zj_match_fallbacks(src, key)
 270         return
 271 
 272     if isinstance(src, dict):
 273         if key in src:
 274             yield key
 275             return
 276 
 277         low = key.casefold()
 278         for k in src.keys():
 279             if low == k.casefold():
 280                 yield k
 281                 return
 282 
 283         yield from zj_match_indices(src, key)
 284         yield from zj_match_fallbacks(src, key)
 285         return
 286 
 287     yield from zj_match_fallbacks(src, key)
 288 
 289 
 290 def zj_match_indices(src: Any, key: str) -> Iterable:
 291     try:
 292         i = int(key)
 293 
 294         if isinstance(src, (list, tuple)):
 295             l = len(src)
 296             yield src[i] if -l <= i < l else None
 297             return
 298 
 299         if isinstance(src, dict):
 300             l = len(src)
 301             if i < 0:
 302                 i += l
 303             if i < 0 or i >= l:
 304                 return
 305 
 306             for j, k in enumerate(src.keys()):
 307                 if i == j:
 308                     yield k
 309                     return
 310 
 311         return
 312     except Exception:
 313         pass
 314 
 315     m = zj_slice_re.match(key)
 316     if not m:
 317         return
 318 
 319     l = len(src)
 320 
 321     (start, _, kind, stop, _) = m.groups()
 322     start = int(start) if start != '' else 0
 323     stop = int(stop) if stop != '' else l
 324 
 325     if start < 0:
 326         start += l
 327     start = max(start, 0)
 328     if stop < 0:
 329         stop += l
 330     stop = min(stop, l)
 331     if kind == '..':
 332         stop += 1
 333     stop = min(stop, l)
 334 
 335     if start > stop:
 336         return
 337     if (start < 0 and stop < 0) or (start >= l and stop >= l):
 338         return
 339 
 340 
 341     if isinstance(src, dict):
 342         for i, k in enumerate(src.keys()):
 343             if i >= stop:
 344                 return
 345             if start <= i:
 346                 yield k
 347         return
 348 
 349     if isinstance(src, (list, tuple)):
 350         yield from range(start, stop)
 351         return
 352 
 353 
 354 
 355 def zj_match_fallbacks(src: Any, key: str) -> Iterable:
 356     fn = zj_fallbacks.get(key, None)
 357     if fn:
 358         yield fn(src)
 359 
 360 
 361 def zj_help(*_) -> NoReturn:
 362     print(zj_info_msg.strip(), file=stderr)
 363     exit(1)
 364 
 365 
 366 def zj_keys(src: Any) -> Any:
 367     if isinstance(src, dict):
 368         return tuple(src.keys())
 369     if isinstance(src, (list, tuple)):
 370         return tuple(range(len(src)))
 371     return None
 372 
 373 
 374 def zj_info(x: Any) -> str:
 375     if isinstance(x, dict):
 376         return f'object ({len(x)} items)'
 377     if isinstance(x, (list, tuple)):
 378         return f'array ({len(x)} items)'
 379     return zj_type(x)
 380 
 381 
 382 def zj_type(x: Any) -> str:
 383     return {
 384         type(None): 'null',
 385         dict: 'object',
 386         float: 'number',
 387         int: 'number',
 388         str: 'string',
 389         list: 'array',
 390         tuple: 'array',
 391     }.get(type(x), 'other')
 392 
 393 
 394 zj_fallbacks: Dict[str, Callable] = {
 395     '.h': zj_help,
 396     '.help': zj_help,
 397     ':h': zj_help,
 398     ':help': zj_help,
 399     ':help:': zj_help,
 400     '.i': zj_info,
 401     '.info': zj_info,
 402     ':i': zj_info,
 403     ':info': zj_info,
 404     ':info:': zj_info,
 405     '.k': zj_keys,
 406     '.keys': zj_keys,
 407     ':keys': zj_keys,
 408     ':keys:': zj_keys,
 409     '.l': len,
 410     '.len': len,
 411     '.length': len,
 412     ':l': len,
 413     ':len': len,
 414     ':length': len,
 415     ':len:': len,
 416     ':length:': len,
 417     '.t': zj_type,
 418     '.type': zj_type,
 419     ':t': zj_type,
 420     ':type': zj_type,
 421     ':type:': zj_type,
 422 }
 423 
 424 
 425 def zj_pick(src: Any, keys: Tuple[str, ...]) -> Any:
 426     if isinstance(src, dict):
 427         picked = {}
 428         for k in keys:
 429             for k in zj_match_keys(src, k):
 430                 picked[k] = src[k]
 431         return picked
 432 
 433     if isinstance(src, (list, tuple)):
 434         picked = []
 435         for k in keys:
 436             for i in zj_match_indices(src, k):
 437                 picked.append(src[i])
 438         return tuple(picked)
 439 
 440     msg = f'can\'t pick properties from value of type {zj_type(src)}'
 441     raise Exception(msg)
 442 
 443 
 444 def zj_drop(src: Any, keys: Tuple[str, ...]) -> Any:
 445     if isinstance(src, dict):
 446         avoid = set()
 447         for k in keys:
 448             for k in zj_match_keys(src, k):
 449                 avoid.add(k)
 450         return {k: v for k, v in src.items() if not k in avoid}
 451 
 452     if isinstance(src, (list, tuple)):
 453         l = len(src)
 454         avoid = set()
 455         for k in keys:
 456             for i in zj_match_indices(src, k):
 457                 avoid.add(i if i >= 0 else i + l)
 458         return tuple(v for i, v in enumerate(src) if not i in avoid)
 459 
 460     msg = f'can\'t drop properties from value of type {zj_type(src)}'
 461     raise Exception(msg)
 462 
 463 
 464 zj_final_fallbacks: Dict[str, Callable] = {
 465     '+': zj_pick,
 466     ':+:': zj_pick,
 467     ':+': zj_pick,
 468     '+:': zj_pick,
 469     '/+': zj_pick,
 470     '+/': zj_pick,
 471 
 472     '-': zj_drop,
 473     ':-:': zj_drop,
 474     ':-': zj_drop,
 475     '-:': zj_drop,
 476     '/-': zj_drop,
 477     '-/': zj_drop,
 478 }
 479 
 480 
 481 apo = '\''
 482 apos = '\''
 483 backquote = '`'
 484 backtick = '`'
 485 ball = ''
 486 block = ''
 487 btick = '`'
 488 bullet = ''
 489 cdot = '·'
 490 circle = ''
 491 cross = '×'
 492 dquo = '"'
 493 dquote = '"'
 494 emdash = ''
 495 endash = ''
 496 ge = ''
 497 geq = ''
 498 hellip = ''
 499 hole = ''
 500 lcurly = '{'
 501 ldquo = ''
 502 ldquote = ''
 503 le = ''
 504 leq = ''
 505 mdash = ''
 506 mdot = '·'
 507 miniball = ''
 508 ndash = ''
 509 neq = ''
 510 rcurly = '}'
 511 rdquo = ''
 512 rdquote = ''
 513 sball = ''
 514 square = ''
 515 squo = '\''
 516 squote = '\''
 517 
 518 
 519 def dive(into: Any, doing: Callable) -> Any:
 520     'Transform a nested value by calling a func via depth-first recursion.'
 521 
 522     # support args in either order
 523     if callable(into):
 524         into, doing = doing, into
 525 
 526     return _dive_kv(None, into, doing)
 527 
 528 deepmap = dive
 529 dive1 = dive
 530 
 531 
 532 def divebin(x: Any, y: Any, doing: Callable) -> Any:
 533     'Nested 2-value version of depth-first-recursive func dive.'
 534 
 535     # support args in either order
 536     if callable(x):
 537         x, y, doing = y, doing, x
 538 
 539     narg = required_arg_count(doing)
 540     if narg == 2:
 541         return dive(x, lambda a: dive(y, lambda b: doing(a, b)))
 542     if narg == 4:
 543         return dive(x, lambda i, a: dive(y, lambda j, b: doing(i, a, j, b)))
 544     raise Exception('divebin(...) only supports funcs with 2 or 4 args')
 545 
 546 bindive = divebin
 547 # diveboth = divebin
 548 # dualdive = divebin
 549 # duodive = divebin
 550 dive2 = divebin
 551 
 552 
 553 def _dive_kv(key: Any, into: Any, doing: Callable) -> Any:
 554     if isinstance(into, dict):
 555         return {k: _dive_kv(k, v, doing) for k, v in into.items()}
 556     if isinstance(into, Iterable) and not isinstance(into, str):
 557         return [_dive_kv(i, e, doing) for i, e in enumerate(into)]
 558 
 559     narg = required_arg_count(doing)
 560     return doing(key, into) if narg == 2 else doing(into)
 561 
 562 
 563 def recover(*args) -> Any:
 564     '''
 565     Catch exceptions using a lambda/callback func, in one of 6 ways
 566         recover(zero_args_func)
 567         recover(zero_args_func, exception_replacement_value)
 568         recover(zero_args_func, one_arg_exception_handling_func)
 569         recover(one_arg_func, arg)
 570         recover(one_arg_func, arg, exception_replacement_value)
 571         recover(one_arg_func, arg, one_arg_exception_handling_func)
 572     '''
 573 
 574     if len(args) == 1:
 575         f = args[0]
 576         try:
 577             return f()
 578         except Exception:
 579             return None
 580     elif len(args) == 2:
 581         f, fallback = args[0], args[1]
 582         if callable(f) and callable(fallback):
 583             try:
 584                 return f()
 585             except Exception as e:
 586                 nargs = required_arg_count(fallback)
 587                 return fallback(e) if nargs == 1 else fallback()
 588         else:
 589             try:
 590                 return f() if required_arg_count(f) == 0 else f(args[1])
 591             except Exception:
 592                 return fallback
 593     elif len(args) == 3:
 594         f, x, fallback = args[0], args[1], args[2]
 595         if callable(f) and callable(fallback):
 596             try:
 597                 return f(x)
 598             except Exception as e:
 599                 nargs = required_arg_count(fallback)
 600                 return fallback(e) if nargs == 1 else fallback()
 601         else:
 602             try:
 603                 return f(x)
 604             except Exception:
 605                 return fallback
 606     else:
 607         raise Exception('recover(...) only works with 1, 2, or 3 args')
 608 
 609 attempt = recover
 610 attempted = recover
 611 recovered = recover
 612 recoverred = recover
 613 rescue = recover
 614 rescued = recover
 615 trycall = recover
 616 
 617 
 618 def required_arg_count(f: Callable) -> int:
 619     if isinstance(f, type):
 620         return 1
 621 
 622     meta = getfullargspec(f)
 623     n = len(meta.args)
 624     if meta.defaults:
 625         n -= len(meta.defaults)
 626     return n
 627 
 628 
 629 typeof = zj_type
 630 
 631 
 632 # deny file-access to expression-evaluators
 633 open = None
 634 
 635 try:
 636     # load data, trying to handle help-like options as well
 637     try:
 638         data = load(stdin.buffer)
 639     except Exception as e:
 640         if len(argv) == 2 and argv[1] in ('-h', '--h', '-help', '--help'):
 641             zj_help(None)
 642         else:
 643             raise e
 644 
 645     data = zj_zoom(data, tuple(argv[1:]))
 646     dump(data, stdout, indent=2, allow_nan=False, check_circular=False)
 647     # dump(data, stdout, indent=None, allow_nan=False, check_circular=False)
 648     stdout.write('\n')
 649 except BrokenPipeError:
 650     # quit quietly, instead of showing a confusing error message
 651     stderr.close()
 652 except KeyboardInterrupt:
 653     exit(2)
 654 except Exception as e:
 655     print(f'\x1b[31m{e}\x1b[0m', file=stderr)
 656     exit(1)