File: tjp.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 tjp [options...] [python expression] [file/URI...]
  28 
  29 
  30 Transform Json with Python runs a python expression on a single JSON-encoded
  31 input.
  32 
  33 The expression can use either `v`, `value`, `d`, or `data` for the decoded
  34 input. Invalid-JSON inputs result in an error, with no chance of recovery.
  35 
  36 Input-sources can be either files or web-URIs. When not given a named input,
  37 the standard input is used.
  38 
  39 
  40 Options
  41 
  42 All options can start with either a single or a double leading dash:
  43 
  44     -c, -compact, -j0, -json0      emit result as compact single-line JSON
  45     -h, -help                      show this help message
  46     -m, -mod, -module, -modules    import modules named in the next argument,
  47                                    where multiple names are comma-separated
  48     -n, -nil, -none, -null         don't read any input
  49     -t, -trace, -traceback         turn exceptions into multi-line tracebacks
  50     -z, -zoom                      zoom into data using rest of the args
  51 
  52 
  53 Examples
  54 
  55 # numbers from 0 to 5; no input is read/used
  56 tjp = 'range(6)'
  57 
  58 # using bases 1 to 5, find all their powers up to the 4th
  59 tjp = '((n**p for p in range(1, 4+1)) for n in range(1, 6))'
  60 
  61 # keep only the last 2 items from the input
  62 tjp = 'range(1, 6)' | tjp 'data[-2:]'
  63 
  64 # chunk/regroup input items into arrays of up to 3 items each
  65 tjp = 'range(1, 8)' | tjp 'chunk(data, 3)'
  66 
  67 # ignore errors/exceptions, in favor of a fallback value
  68 tjp = 'rescue(lambda: 2 * float("no way"), "fallback value")'
  69 
  70 # ignore errors/exceptions, calling a fallback func with the exception
  71 tjp = 'rescue(lambda: 2 * float("no way"), str)'
  72 
  73 # use dot-syntax on JSON data
  74 tjp = '{"abc": {"xyz": 123}}' | tjp -dots 'data.abc.xyz'
  75 
  76 # use dot-syntax on JSON data; keywords as properties are syntax-errors
  77 tjp = '{"abc": {"def": 123}}' | tjp -dots 'data.abc["def"]'
  78 
  79 # func results are automatically called on the input
  80 tjp = '{"abc": 123, "def": 456}' | tjp len
  81 
  82 # an array of 10 random integers between 1 and 10
  83 tjp -m random = '(random.randint(1, 10) for _ in range(10))'
  84 
  85 # try to auto-parse values (esp. numbers) from a table of string values
  86 echo '[{"key": "abc", "val": "123"}, {"key": "xyz", "val": "no"}]' | \\
  87     tjp '[{k: rescue(lambda: loads(v), v) for k, v in e.items()} for e in v]'
  88 
  89 # zoom into last array-item of the value of object under `def` key
  90 echo '{"abc": 123, "def": [456, 789]}' | tjp -z def -1
  91 
  92 # zoom into last array-item of the last value of object
  93 echo '{"abc": 123, "def": [456, 789]}' | tjp -z -1 -1
  94 
  95 # zoom into last value of object
  96 echo '{"abc": 123, "def": [456, 789]}' | tjp -z -1
  97 
  98 # zoom into value of object, using key matched case-insensitively
  99 echo '{"abc": 123, "def": [456, 789]}' | tjp -z ABC
 100 '''
 101 
 102 
 103 from itertools import islice
 104 from json import dump, load, loads
 105 from math import isnan
 106 from re import compile as compile_uncached, IGNORECASE
 107 from sys import argv, exit, stderr, stdin, stdout
 108 from typing import Iterable
 109 
 110 
 111 if len(argv) < 2:
 112     print(info.strip(), file=stderr)
 113     exit(0)
 114 if len(argv) > 1 and argv[1] in ('-h', '--h', '-help', '--help'):
 115     print(info.strip())
 116     exit(0)
 117 
 118 
 119 class Skip:
 120     pass
 121 
 122 skip = Skip()
 123 
 124 class Dottable:
 125     'Enable convenient dot-syntax access to dictionary values.'
 126 
 127     def __getattr__(self, key):
 128         return self.__dict__[key] if key in self.__dict__ else None
 129 
 130     def __getitem__(self, key):
 131         return self.__dict__[key] if key in self.__dict__ else None
 132 
 133     def __iter__(self):
 134         return iter(self.__dict__)
 135 
 136 def dotate(x):
 137     'Recursively ensure all dictionaries in a value are dot-accessible.'
 138 
 139     if isinstance(x, dict):
 140         d = Dottable()
 141         d.__dict__ = {k: dotate(v) for k, v in x.items()}
 142         return d
 143     if isinstance(x, list):
 144         return [dotate(e) for e in x]
 145     if isinstance(x, tuple):
 146         return tuple(dotate(e) for e in x)
 147     return x
 148 
 149 dotated = dote = doted = dotified = dotify = dottified = dottify = dotate
 150 
 151 def chunk(items, chunk_size):
 152     'Break iterable into chunks, each with up to the item-count given.'
 153 
 154     if isinstance(items, str):
 155         n = len(items)
 156         while n >= chunk_size:
 157             yield items[:chunk_size]
 158             items = items[chunk_size:]
 159             n -= chunk_size
 160         if n > 0:
 161             yield items
 162         return
 163 
 164     if not isinstance(chunk_size, int):
 165         raise Exception('non-integer chunk-size')
 166     if chunk_size < 1:
 167         raise Exception('non-positive chunk-size')
 168 
 169     it = iter(items)
 170     while True:
 171         head = tuple(islice(it, chunk_size))
 172         if not head:
 173             return
 174         yield head
 175 
 176 chunked = chunk
 177 
 178 # re_cache is used by custom func compile to cache previously-compiled
 179 # regular-expressions, which makes them quicker to (re)use in formulas
 180 re_cache = {}
 181 
 182 def re_compile(expr, flags = 0):
 183     'Speed-up using regexes, by avoiding recompilations.'
 184 
 185     if flags in re_cache:
 186         cache = re_cache[flags]
 187     else:
 188         cache = {}
 189         re_cache[flags] = cache
 190     if expr in cache:
 191         return cache[expr]
 192 
 193     pat = compile_uncached(expr, flags)
 194     cache[expr] = pat
 195     return pat
 196 
 197 def icompile(expr):
 198     return re_compile(expr, IGNORECASE)
 199 
 200 def cond(*args):
 201     if len(args) == 0:
 202         return None
 203 
 204     for i, e in enumerate(args):
 205         if i % 2 == 0 and i < len(args) - 1 and e:
 206             return args[i + 1]
 207 
 208     return args[-1] if len(args) % 2 == 1 else None
 209 
 210 def dive(into, using):
 211     'Depth-first recursive caller for 1-input functions.'
 212 
 213     if callable(into):
 214         into, using = using, into
 215 
 216     def rec(v):
 217         if isinstance(v, dict):
 218             return {k: rec(v) for k, v in v.items()}
 219         if isinstance(v, Iterable) and not isinstance(v, str):
 220             return [rec(v) for v in v]
 221         return using(v)
 222 
 223     return rec(into)
 224 
 225 def divekeys(into, using):
 226     'Depth-first recursive caller for 2-input funcs which rename dict keys.'
 227 
 228     if callable(into):
 229         into, using = using, into
 230 
 231     def rec(v):
 232         if isinstance(v, dict):
 233             return {using(k): rec(v) for k, v in v.items()}
 234         if isinstance(v, Iterable) and not isinstance(v, str):
 235             return [rec(v) for i, v in enumerate(v)]
 236         return v
 237 
 238     return rec(None, into)
 239 
 240 def divekv(into, using, using2 = None):
 241     'Depth-first recursive caller for 2-input functions.'
 242 
 243     if using2 is None:
 244         if callable(into):
 245             into, using = using, into
 246     else:
 247         if not callable(using2):
 248             into, using, using2 = using2, into, using
 249 
 250     def rec(k, v):
 251         if isinstance(v, dict):
 252             return {k: rec(k, v) for k, v in v.items()}
 253         if isinstance(v, Iterable) and not isinstance(v, str):
 254             return [rec(i, v) for i, v in enumerate(v)]
 255         return using(k, v)
 256 
 257     def rec2(k, v):
 258         if isinstance(v, dict):
 259             return {str(using(k, v)): rec2(k, v) for k, v in v.items()}
 260         if isinstance(v, Iterable) and not isinstance(v, str):
 261             # return {str(using(i, v)): rec2(i, v) for i, v in enumerate(v)}
 262             return [rec2(i, v) for i, v in enumerate(v)]
 263         return using2(k, v)
 264 
 265     return rec(None, into) if using2 is None else rec2(None, into)
 266 
 267 kvdive = divekv
 268 
 269 def drop(src, *what):
 270     if isinstance(src, str):
 271         for s in what:
 272             src = src.replace(s, '')
 273         return src
 274 
 275     def kdrop(src, what):
 276         return {k: v for (k, v) in src.items() if not (k in what)}
 277 
 278     if isinstance(src, dict):
 279         return kdrop(src, set(what))
 280 
 281     if isinstance(src, Iterable):
 282         what = set(what)
 283         return [kdrop(e, what) for e in src if isinstance(e, dict)]
 284 
 285     return None
 286 
 287 dropped = drop
 288 
 289 def join(x, y = ' '):
 290     'Join values into a string, or make a dict from keys and values.'
 291 
 292     if isinstance(x, str):
 293         return x.join(str(v) for v in y)
 294     if isinstance(y, str):
 295         return y.join(str(v) for v in x)
 296     return {k: v for k, v in zip(x, y)}
 297 
 298 def maybe(f, x):
 299     try:
 300         return f(x)
 301     except Exception as _:
 302         return x
 303 
 304 def number(x):
 305     try:
 306         return int(x)
 307     except Exception as _:
 308         pass
 309     try:
 310         return float(x)
 311     except Exception as _:
 312         return x
 313 
 314 def pick(src, *keys):
 315     if isinstance(src, dict):
 316         return {k: src.get(k, None) for k in keys}
 317     return [{k: e.get(k, None) for k in keys} for e in src if isinstance(e, dict)]
 318 
 319 picked = pick
 320 
 321 def rescue(attempt, fallback = None):
 322     try:
 323         return attempt()
 324     except BrokenPipeError as e:
 325         raise e
 326     except Exception as e:
 327         if callable(fallback):
 328             return fallback(e)
 329         return fallback
 330 
 331 rescued = rescue
 332 
 333 def retype(x):
 334     'Try to narrow the type of the value given.'
 335 
 336     if isinstance(x, float):
 337         n = int(x)
 338         return n if float(n) == x else x
 339 
 340     if not isinstance(x, str):
 341         return x
 342 
 343     try:
 344         return loads(x)
 345     except Exception:
 346         pass
 347 
 348     try:
 349         return int(x)
 350     except Exception:
 351         pass
 352 
 353     try:
 354         return float(x)
 355     except Exception:
 356         pass
 357 
 358     return x
 359 
 360 autocast = autocasted = mold = molded = recast = recasted = remold = retype
 361 remolded = retyped = retype
 362 
 363 def typeof(x):
 364     # return str(type(x))
 365     return {
 366         type(None): 'null',
 367         bool: 'boolean',
 368         dict: 'object',
 369         float: 'number',
 370         int: 'number',
 371         str: 'string',
 372         list: 'array',
 373         tuple: 'array',
 374     }.get(type(x), 'other')
 375 
 376 jstype = typeof
 377 
 378 
 379 def result_needs_fixing(x):
 380     if isinstance(x, float):
 381         return not isnan(x)
 382     if x is None or isinstance(x, (bool, int, float, str)):
 383         return False
 384     rec = result_needs_fixing
 385     if isinstance(x, dict):
 386         return any(rec(k) or rec(v) for k, v in x.items())
 387     if isinstance(x, (list, tuple)):
 388         return any(rec(e) for e in x)
 389     return True
 390 
 391 def fix_result(x, default):
 392     if x is type:
 393         return type(default).__name__
 394 
 395     # if expression results in a func, auto-call it with the original data
 396     if callable(x):
 397         x = x(default)
 398 
 399     if isinstance(x, float) and isnan(x):
 400         return None
 401 
 402     if x is None or isinstance(x, (bool, int, float, str)):
 403         return x
 404 
 405     rec = fix_result
 406 
 407     if isinstance(x, dict):
 408         return {
 409             rec(k, default): rec(v, default) for k, v in x.items() if not
 410                 (isinstance(k, Skip) or isinstance(v, Skip))
 411         }
 412 
 413     if isinstance(x, Iterable):
 414         return tuple(rec(e, default) for e in x if not isinstance(e, Skip))
 415 
 416     if isinstance(x, Dottable):
 417         return rec(x.__dict__, default)
 418 
 419     if isinstance(x, Exception):
 420         raise x
 421 
 422     return None if isinstance(x, Skip) else str(x)
 423 
 424 def fail(msg, code = 1):
 425     print(str(msg), file=stderr)
 426     exit(code)
 427 
 428 def message(msg, result = None):
 429     print(msg, file=stderr)
 430     return result
 431 
 432 msg = message
 433 
 434 def seemsurl(path):
 435     protocols = ('https://', 'http://', 'file://', 'ftp://', 'data:')
 436     return any(path.startswith(p) for p in protocols)
 437 
 438 def matchkey(kv, key):
 439     if key in kv:
 440         return key
 441 
 442     low = key.lower()
 443     for k in kv.keys():
 444         if low == k.lower():
 445             return k
 446 
 447     try:
 448         i = int(key)
 449         l = len(kv)
 450         if i < 0:
 451             i += l
 452 
 453         if not (-l <= i < l):
 454             return key
 455 
 456         for j, k in enumerate(kv.keys()):
 457             if i == j:
 458                 return k
 459     except Exception:
 460         return key
 461 
 462     return key
 463 
 464 def zoom(data, keys):
 465     for i, k in enumerate(keys):
 466         if isinstance(data, dict):
 467             # m = matchkey(data, k)
 468             # if not (m in data):
 469             #     raise Exception(f'{m}: object doesn\'t have that key')
 470             data = data.get(matchkey(data, k), None)
 471             continue
 472 
 473         if isinstance(data, (list, tuple)):
 474             if k == '+':
 475                 pick = keys[i + 1:]
 476                 return [{k: e.get(k, None) for k in pick}
 477                         for e in data if isinstance(e, dict)]
 478             if k == '-':
 479                 avoid = set(keys[i + 1:])
 480                 return [{k: v for (k, v) in e.items() if not (k in avoid)}
 481                         for e in data if isinstance(e, dict)]
 482             if k == '.':
 483                 rest = keys[i + 1:]
 484                 return [zoom(e, rest) for e in data]
 485 
 486             try:
 487                 k = int(k)
 488                 l = len(data)
 489                 data = data[k] if -l <= k < l else None
 490             except Exception:
 491                 # raise Exception(f'{k}: arrays don\'t have keys like objects')
 492                 data = None
 493             continue
 494 
 495         # return None
 496         # data = None
 497         raise Exception(f'{k}: can\'t zoom on value of type {typeof(data)}')
 498 
 499     return data
 500 
 501 def make_eval_once(run):
 502     def eval_once(expr):
 503         global eval
 504         eval = None
 505         return run(expr)
 506     return eval_once
 507 
 508 
 509 cr = '\r'
 510 crlf = '\r\n'
 511 dquo = dquote = '"'
 512 empty = ''
 513 lcurly = '{'
 514 lf = '\n'
 515 rcurly = '}'
 516 squo = squote = '\''
 517 
 518 nil = none = null = None
 519 
 520 
 521 no_input_opts = (
 522     '=', '-n', '--n', '-nil', '--nil', '-none', '--none', '-null', '--null',
 523 )
 524 compact_output_opts = (
 525     '-c', '--c', '-compact', '--compact', '-j0', '--j0', '-json0', '--json0',
 526 )
 527 dot_opts = ('--d', '-dot', '--dot', '-dots', '--dots')
 528 modules_opts = (
 529     '-m', '--m', '-mod', '--mod', '-module', '--module',
 530     '-modules', '--modules',
 531 )
 532 pipe_opts = ('-p', '--p', '-pipe', '--pipe')
 533 trace_opts = ('-t', '--t', '-trace', '--trace', '-traceback', '--traceback')
 534 zoom_opts = ('-z', '--z', '-zoom', '--zoom')
 535 
 536 args = argv[1:]
 537 no_input = False
 538 zoom_stdin = False
 539 pipe_mode = False
 540 trace_errors = False
 541 dottable_input = False
 542 compact_output = False
 543 
 544 while len(args) > 0:
 545     if args[0] == '--':
 546         args = args[1:]
 547         break
 548 
 549     if args[0] in no_input_opts:
 550         no_input = True
 551         args = args[1:]
 552         continue
 553 
 554     if args[0] in compact_output_opts:
 555         compact_output = True
 556         args = args[1:]
 557         continue
 558 
 559     if args[0] in dot_opts:
 560         dottable_input = True
 561         args = args[1:]
 562         continue
 563 
 564     if args[0] in pipe_opts:
 565         pipe_mode = True
 566         args = args[1:]
 567         break
 568 
 569     if args[0] in modules_opts:
 570         try:
 571             if len(args) < 2:
 572                 msg = 'a module name or a comma-separated list of modules'
 573                 raise Exception('expected ' + msg)
 574 
 575             g = globals()
 576             from importlib import import_module
 577             for e in args[1].split(','):
 578                 g[e] = import_module(e)
 579 
 580             g = None
 581             import_module = None
 582             args = args[2:]
 583         except Exception as e:
 584             fail(e, 1)
 585 
 586         continue
 587 
 588     if args[0] in trace_opts:
 589         trace_errors = True
 590         args = args[1:]
 591         continue
 592 
 593     if args[0] in zoom_opts:
 594         zoom_stdin = True
 595         args = args[1:]
 596         break
 597 
 598     break
 599 
 600 
 601 try:
 602     if zoom_stdin:
 603         data = load(stdin)
 604         data = zoom(data, args)
 605         v = data
 606     else:
 607         expr = 'data'
 608         if len(args) > 0 and (not pipe_mode):
 609             expr = args[0]
 610             args = args[1:]
 611 
 612         if expr == '.':
 613             expr = 'data'
 614         if not pipe_mode:
 615             expr = compile(expr, expr, mode='eval')
 616 
 617         if (not pipe_mode) and len(args) > 1:
 618             raise Exception('can\'t use more than 1 input')
 619         path = '-' if len(args) == 0 or pipe_mode else args[0]
 620 
 621         if no_input:
 622             data = None
 623         elif path == '-':
 624             data = load(stdin)
 625         elif seemsurl(path):
 626             from io import TextIOWrapper
 627             from urllib.request import urlopen
 628             with urlopen(path) as inp:
 629                 with TextIOWrapper(inp, encoding='utf-8') as txt:
 630                     data = load(txt)
 631         else:
 632             with open(path, encoding='utf-8') as inp:
 633                 data = load(inp)
 634 
 635         if dottable_input:
 636             data = dotate(data)
 637 
 638         v = val = value = d = dat = data
 639         exec = None
 640         open = None
 641         compile = None
 642 
 643         if pipe_mode:
 644             funcs = [eval(s) for s in args]
 645             eval = None
 646 
 647             # variable names `o` and `p` work like in the `pyp` tool, except
 648             # the pipeline steps were given as separate cmd-line arguments
 649             global o, p
 650 
 651             o = p = prev = v
 652             for f in funcs:
 653                 p = f(p)
 654                 if callable(p):
 655                     p = p(prev)
 656                 prev = p
 657             v = p
 658         else:
 659             eval = make_eval_once(eval)
 660             v = eval(expr)
 661 
 662     if result_needs_fixing(v):
 663         v = fix_result(v, value)
 664 
 665     if compact_output:
 666         dump(v, stdout, indent=None, separators=(',', ':'), allow_nan=False)
 667     else:
 668         dump(v, stdout, indent=2, separators=(',', ': '), allow_nan=False)
 669     print()
 670 except BrokenPipeError:
 671     # quit quietly, instead of showing a confusing error message
 672     stderr.close()
 673     exit(0)
 674 except KeyboardInterrupt:
 675     exit(2)
 676 except Exception as e:
 677     if trace_errors:
 678         raise e
 679     else:
 680         fail(e, 1)