functions.go (18532B)
1 package jmespath 2 3 import ( 4 "encoding/json" 5 "errors" 6 "fmt" 7 "math" 8 "reflect" 9 "sort" 10 "strconv" 11 "strings" 12 "unicode/utf8" 13 ) 14 15 type jpFunction func(arguments []interface{}) (interface{}, error) 16 17 type jpType string 18 19 const ( 20 jpUnknown jpType = "unknown" 21 jpNumber jpType = "number" 22 jpString jpType = "string" 23 jpArray jpType = "array" 24 jpObject jpType = "object" 25 jpArrayNumber jpType = "array[number]" 26 jpArrayString jpType = "array[string]" 27 jpExpref jpType = "expref" 28 jpAny jpType = "any" 29 ) 30 31 type functionEntry struct { 32 name string 33 arguments []argSpec 34 handler jpFunction 35 hasExpRef bool 36 } 37 38 type argSpec struct { 39 types []jpType 40 variadic bool 41 } 42 43 type byExprString struct { 44 intr *treeInterpreter 45 node ASTNode 46 items []interface{} 47 hasError bool 48 } 49 50 func (a *byExprString) Len() int { 51 return len(a.items) 52 } 53 func (a *byExprString) Swap(i, j int) { 54 a.items[i], a.items[j] = a.items[j], a.items[i] 55 } 56 func (a *byExprString) Less(i, j int) bool { 57 first, err := a.intr.Execute(a.node, a.items[i]) 58 if err != nil { 59 a.hasError = true 60 // Return a dummy value. 61 return true 62 } 63 ith, ok := first.(string) 64 if !ok { 65 a.hasError = true 66 return true 67 } 68 second, err := a.intr.Execute(a.node, a.items[j]) 69 if err != nil { 70 a.hasError = true 71 // Return a dummy value. 72 return true 73 } 74 jth, ok := second.(string) 75 if !ok { 76 a.hasError = true 77 return true 78 } 79 return ith < jth 80 } 81 82 type byExprFloat struct { 83 intr *treeInterpreter 84 node ASTNode 85 items []interface{} 86 hasError bool 87 } 88 89 func (a *byExprFloat) Len() int { 90 return len(a.items) 91 } 92 func (a *byExprFloat) Swap(i, j int) { 93 a.items[i], a.items[j] = a.items[j], a.items[i] 94 } 95 func (a *byExprFloat) Less(i, j int) bool { 96 first, err := a.intr.Execute(a.node, a.items[i]) 97 if err != nil { 98 a.hasError = true 99 // Return a dummy value. 100 return true 101 } 102 ith, ok := first.(float64) 103 if !ok { 104 a.hasError = true 105 return true 106 } 107 second, err := a.intr.Execute(a.node, a.items[j]) 108 if err != nil { 109 a.hasError = true 110 // Return a dummy value. 111 return true 112 } 113 jth, ok := second.(float64) 114 if !ok { 115 a.hasError = true 116 return true 117 } 118 return ith < jth 119 } 120 121 type functionCaller struct { 122 functionTable map[string]functionEntry 123 } 124 125 func newFunctionCaller() *functionCaller { 126 caller := &functionCaller{} 127 caller.functionTable = map[string]functionEntry{ 128 "length": { 129 name: "length", 130 arguments: []argSpec{ 131 {types: []jpType{jpString, jpArray, jpObject}}, 132 }, 133 handler: jpfLength, 134 }, 135 "starts_with": { 136 name: "starts_with", 137 arguments: []argSpec{ 138 {types: []jpType{jpString}}, 139 {types: []jpType{jpString}}, 140 }, 141 handler: jpfStartsWith, 142 }, 143 "abs": { 144 name: "abs", 145 arguments: []argSpec{ 146 {types: []jpType{jpNumber}}, 147 }, 148 handler: jpfAbs, 149 }, 150 "avg": { 151 name: "avg", 152 arguments: []argSpec{ 153 {types: []jpType{jpArrayNumber}}, 154 }, 155 handler: jpfAvg, 156 }, 157 "ceil": { 158 name: "ceil", 159 arguments: []argSpec{ 160 {types: []jpType{jpNumber}}, 161 }, 162 handler: jpfCeil, 163 }, 164 "contains": { 165 name: "contains", 166 arguments: []argSpec{ 167 {types: []jpType{jpArray, jpString}}, 168 {types: []jpType{jpAny}}, 169 }, 170 handler: jpfContains, 171 }, 172 "ends_with": { 173 name: "ends_with", 174 arguments: []argSpec{ 175 {types: []jpType{jpString}}, 176 {types: []jpType{jpString}}, 177 }, 178 handler: jpfEndsWith, 179 }, 180 "floor": { 181 name: "floor", 182 arguments: []argSpec{ 183 {types: []jpType{jpNumber}}, 184 }, 185 handler: jpfFloor, 186 }, 187 "map": { 188 name: "amp", 189 arguments: []argSpec{ 190 {types: []jpType{jpExpref}}, 191 {types: []jpType{jpArray}}, 192 }, 193 handler: jpfMap, 194 hasExpRef: true, 195 }, 196 "max": { 197 name: "max", 198 arguments: []argSpec{ 199 {types: []jpType{jpArrayNumber, jpArrayString}}, 200 }, 201 handler: jpfMax, 202 }, 203 "merge": { 204 name: "merge", 205 arguments: []argSpec{ 206 {types: []jpType{jpObject}, variadic: true}, 207 }, 208 handler: jpfMerge, 209 }, 210 "max_by": { 211 name: "max_by", 212 arguments: []argSpec{ 213 {types: []jpType{jpArray}}, 214 {types: []jpType{jpExpref}}, 215 }, 216 handler: jpfMaxBy, 217 hasExpRef: true, 218 }, 219 "sum": { 220 name: "sum", 221 arguments: []argSpec{ 222 {types: []jpType{jpArrayNumber}}, 223 }, 224 handler: jpfSum, 225 }, 226 "min": { 227 name: "min", 228 arguments: []argSpec{ 229 {types: []jpType{jpArrayNumber, jpArrayString}}, 230 }, 231 handler: jpfMin, 232 }, 233 "min_by": { 234 name: "min_by", 235 arguments: []argSpec{ 236 {types: []jpType{jpArray}}, 237 {types: []jpType{jpExpref}}, 238 }, 239 handler: jpfMinBy, 240 hasExpRef: true, 241 }, 242 "type": { 243 name: "type", 244 arguments: []argSpec{ 245 {types: []jpType{jpAny}}, 246 }, 247 handler: jpfType, 248 }, 249 "keys": { 250 name: "keys", 251 arguments: []argSpec{ 252 {types: []jpType{jpObject}}, 253 }, 254 handler: jpfKeys, 255 }, 256 "values": { 257 name: "values", 258 arguments: []argSpec{ 259 {types: []jpType{jpObject}}, 260 }, 261 handler: jpfValues, 262 }, 263 "sort": { 264 name: "sort", 265 arguments: []argSpec{ 266 {types: []jpType{jpArrayString, jpArrayNumber}}, 267 }, 268 handler: jpfSort, 269 }, 270 "sort_by": { 271 name: "sort_by", 272 arguments: []argSpec{ 273 {types: []jpType{jpArray}}, 274 {types: []jpType{jpExpref}}, 275 }, 276 handler: jpfSortBy, 277 hasExpRef: true, 278 }, 279 "join": { 280 name: "join", 281 arguments: []argSpec{ 282 {types: []jpType{jpString}}, 283 {types: []jpType{jpArrayString}}, 284 }, 285 handler: jpfJoin, 286 }, 287 "reverse": { 288 name: "reverse", 289 arguments: []argSpec{ 290 {types: []jpType{jpArray, jpString}}, 291 }, 292 handler: jpfReverse, 293 }, 294 "to_array": { 295 name: "to_array", 296 arguments: []argSpec{ 297 {types: []jpType{jpAny}}, 298 }, 299 handler: jpfToArray, 300 }, 301 "to_string": { 302 name: "to_string", 303 arguments: []argSpec{ 304 {types: []jpType{jpAny}}, 305 }, 306 handler: jpfToString, 307 }, 308 "to_number": { 309 name: "to_number", 310 arguments: []argSpec{ 311 {types: []jpType{jpAny}}, 312 }, 313 handler: jpfToNumber, 314 }, 315 "not_null": { 316 name: "not_null", 317 arguments: []argSpec{ 318 {types: []jpType{jpAny}, variadic: true}, 319 }, 320 handler: jpfNotNull, 321 }, 322 } 323 return caller 324 } 325 326 func (e *functionEntry) resolveArgs(arguments []interface{}) ([]interface{}, error) { 327 if len(e.arguments) == 0 { 328 return arguments, nil 329 } 330 if !e.arguments[len(e.arguments)-1].variadic { 331 if len(e.arguments) != len(arguments) { 332 return nil, errors.New("incorrect number of args") 333 } 334 for i, spec := range e.arguments { 335 userArg := arguments[i] 336 err := spec.typeCheck(userArg) 337 if err != nil { 338 return nil, err 339 } 340 } 341 return arguments, nil 342 } 343 if len(arguments) < len(e.arguments) { 344 return nil, errors.New("Invalid arity.") 345 } 346 return arguments, nil 347 } 348 349 func (a *argSpec) typeCheck(arg interface{}) error { 350 for _, t := range a.types { 351 switch t { 352 case jpNumber: 353 if _, ok := arg.(float64); ok { 354 return nil 355 } 356 case jpString: 357 if _, ok := arg.(string); ok { 358 return nil 359 } 360 case jpArray: 361 if isSliceType(arg) { 362 return nil 363 } 364 case jpObject: 365 if _, ok := arg.(map[string]interface{}); ok { 366 return nil 367 } 368 case jpArrayNumber: 369 if _, ok := toArrayNum(arg); ok { 370 return nil 371 } 372 case jpArrayString: 373 if _, ok := toArrayStr(arg); ok { 374 return nil 375 } 376 case jpAny: 377 return nil 378 case jpExpref: 379 if _, ok := arg.(expRef); ok { 380 return nil 381 } 382 } 383 } 384 return fmt.Errorf("Invalid type for: %v, expected: %#v", arg, a.types) 385 } 386 387 func (f *functionCaller) CallFunction(name string, arguments []interface{}, intr *treeInterpreter) (interface{}, error) { 388 entry, ok := f.functionTable[name] 389 if !ok { 390 return nil, errors.New("unknown function: " + name) 391 } 392 resolvedArgs, err := entry.resolveArgs(arguments) 393 if err != nil { 394 return nil, err 395 } 396 if entry.hasExpRef { 397 var extra []interface{} 398 extra = append(extra, intr) 399 resolvedArgs = append(extra, resolvedArgs...) 400 } 401 return entry.handler(resolvedArgs) 402 } 403 404 func jpfAbs(arguments []interface{}) (interface{}, error) { 405 num := arguments[0].(float64) 406 return math.Abs(num), nil 407 } 408 409 func jpfLength(arguments []interface{}) (interface{}, error) { 410 arg := arguments[0] 411 if c, ok := arg.(string); ok { 412 return float64(utf8.RuneCountInString(c)), nil 413 } else if isSliceType(arg) { 414 v := reflect.ValueOf(arg) 415 return float64(v.Len()), nil 416 } else if c, ok := arg.(map[string]interface{}); ok { 417 return float64(len(c)), nil 418 } 419 return nil, errors.New("could not compute length()") 420 } 421 422 func jpfStartsWith(arguments []interface{}) (interface{}, error) { 423 search := arguments[0].(string) 424 prefix := arguments[1].(string) 425 return strings.HasPrefix(search, prefix), nil 426 } 427 428 func jpfAvg(arguments []interface{}) (interface{}, error) { 429 // We've already type checked the value so we can safely use 430 // type assertions. 431 args := arguments[0].([]interface{}) 432 length := float64(len(args)) 433 numerator := 0.0 434 for _, n := range args { 435 numerator += n.(float64) 436 } 437 return numerator / length, nil 438 } 439 func jpfCeil(arguments []interface{}) (interface{}, error) { 440 val := arguments[0].(float64) 441 return math.Ceil(val), nil 442 } 443 func jpfContains(arguments []interface{}) (interface{}, error) { 444 search := arguments[0] 445 el := arguments[1] 446 if searchStr, ok := search.(string); ok { 447 if elStr, ok := el.(string); ok { 448 return strings.Index(searchStr, elStr) != -1, nil 449 } 450 return false, nil 451 } 452 // Otherwise this is a generic contains for []interface{} 453 general := search.([]interface{}) 454 for _, item := range general { 455 if item == el { 456 return true, nil 457 } 458 } 459 return false, nil 460 } 461 func jpfEndsWith(arguments []interface{}) (interface{}, error) { 462 search := arguments[0].(string) 463 suffix := arguments[1].(string) 464 return strings.HasSuffix(search, suffix), nil 465 } 466 func jpfFloor(arguments []interface{}) (interface{}, error) { 467 val := arguments[0].(float64) 468 return math.Floor(val), nil 469 } 470 func jpfMap(arguments []interface{}) (interface{}, error) { 471 intr := arguments[0].(*treeInterpreter) 472 exp := arguments[1].(expRef) 473 node := exp.ref 474 arr := arguments[2].([]interface{}) 475 mapped := make([]interface{}, 0, len(arr)) 476 for _, value := range arr { 477 current, err := intr.Execute(node, value) 478 if err != nil { 479 return nil, err 480 } 481 mapped = append(mapped, current) 482 } 483 return mapped, nil 484 } 485 func jpfMax(arguments []interface{}) (interface{}, error) { 486 if items, ok := toArrayNum(arguments[0]); ok { 487 if len(items) == 0 { 488 return nil, nil 489 } 490 if len(items) == 1 { 491 return items[0], nil 492 } 493 best := items[0] 494 for _, item := range items[1:] { 495 if item > best { 496 best = item 497 } 498 } 499 return best, nil 500 } 501 // Otherwise we're dealing with a max() of strings. 502 items, _ := toArrayStr(arguments[0]) 503 if len(items) == 0 { 504 return nil, nil 505 } 506 if len(items) == 1 { 507 return items[0], nil 508 } 509 best := items[0] 510 for _, item := range items[1:] { 511 if item > best { 512 best = item 513 } 514 } 515 return best, nil 516 } 517 func jpfMerge(arguments []interface{}) (interface{}, error) { 518 final := make(map[string]interface{}) 519 for _, m := range arguments { 520 mapped := m.(map[string]interface{}) 521 for key, value := range mapped { 522 final[key] = value 523 } 524 } 525 return final, nil 526 } 527 func jpfMaxBy(arguments []interface{}) (interface{}, error) { 528 intr := arguments[0].(*treeInterpreter) 529 arr := arguments[1].([]interface{}) 530 exp := arguments[2].(expRef) 531 node := exp.ref 532 if len(arr) == 0 { 533 return nil, nil 534 } else if len(arr) == 1 { 535 return arr[0], nil 536 } 537 start, err := intr.Execute(node, arr[0]) 538 if err != nil { 539 return nil, err 540 } 541 switch t := start.(type) { 542 case float64: 543 bestVal := t 544 bestItem := arr[0] 545 for _, item := range arr[1:] { 546 result, err := intr.Execute(node, item) 547 if err != nil { 548 return nil, err 549 } 550 current, ok := result.(float64) 551 if !ok { 552 return nil, errors.New("invalid type, must be number") 553 } 554 if current > bestVal { 555 bestVal = current 556 bestItem = item 557 } 558 } 559 return bestItem, nil 560 case string: 561 bestVal := t 562 bestItem := arr[0] 563 for _, item := range arr[1:] { 564 result, err := intr.Execute(node, item) 565 if err != nil { 566 return nil, err 567 } 568 current, ok := result.(string) 569 if !ok { 570 return nil, errors.New("invalid type, must be string") 571 } 572 if current > bestVal { 573 bestVal = current 574 bestItem = item 575 } 576 } 577 return bestItem, nil 578 default: 579 return nil, errors.New("invalid type, must be number of string") 580 } 581 } 582 func jpfSum(arguments []interface{}) (interface{}, error) { 583 items, _ := toArrayNum(arguments[0]) 584 sum := 0.0 585 for _, item := range items { 586 sum += item 587 } 588 return sum, nil 589 } 590 591 func jpfMin(arguments []interface{}) (interface{}, error) { 592 if items, ok := toArrayNum(arguments[0]); ok { 593 if len(items) == 0 { 594 return nil, nil 595 } 596 if len(items) == 1 { 597 return items[0], nil 598 } 599 best := items[0] 600 for _, item := range items[1:] { 601 if item < best { 602 best = item 603 } 604 } 605 return best, nil 606 } 607 items, _ := toArrayStr(arguments[0]) 608 if len(items) == 0 { 609 return nil, nil 610 } 611 if len(items) == 1 { 612 return items[0], nil 613 } 614 best := items[0] 615 for _, item := range items[1:] { 616 if item < best { 617 best = item 618 } 619 } 620 return best, nil 621 } 622 623 func jpfMinBy(arguments []interface{}) (interface{}, error) { 624 intr := arguments[0].(*treeInterpreter) 625 arr := arguments[1].([]interface{}) 626 exp := arguments[2].(expRef) 627 node := exp.ref 628 if len(arr) == 0 { 629 return nil, nil 630 } else if len(arr) == 1 { 631 return arr[0], nil 632 } 633 start, err := intr.Execute(node, arr[0]) 634 if err != nil { 635 return nil, err 636 } 637 if t, ok := start.(float64); ok { 638 bestVal := t 639 bestItem := arr[0] 640 for _, item := range arr[1:] { 641 result, err := intr.Execute(node, item) 642 if err != nil { 643 return nil, err 644 } 645 current, ok := result.(float64) 646 if !ok { 647 return nil, errors.New("invalid type, must be number") 648 } 649 if current < bestVal { 650 bestVal = current 651 bestItem = item 652 } 653 } 654 return bestItem, nil 655 } else if t, ok := start.(string); ok { 656 bestVal := t 657 bestItem := arr[0] 658 for _, item := range arr[1:] { 659 result, err := intr.Execute(node, item) 660 if err != nil { 661 return nil, err 662 } 663 current, ok := result.(string) 664 if !ok { 665 return nil, errors.New("invalid type, must be string") 666 } 667 if current < bestVal { 668 bestVal = current 669 bestItem = item 670 } 671 } 672 return bestItem, nil 673 } else { 674 return nil, errors.New("invalid type, must be number of string") 675 } 676 } 677 func jpfType(arguments []interface{}) (interface{}, error) { 678 arg := arguments[0] 679 if _, ok := arg.(float64); ok { 680 return "number", nil 681 } 682 if _, ok := arg.(string); ok { 683 return "string", nil 684 } 685 if _, ok := arg.([]interface{}); ok { 686 return "array", nil 687 } 688 if _, ok := arg.(map[string]interface{}); ok { 689 return "object", nil 690 } 691 if arg == nil { 692 return "null", nil 693 } 694 if arg == true || arg == false { 695 return "boolean", nil 696 } 697 return nil, errors.New("unknown type") 698 } 699 func jpfKeys(arguments []interface{}) (interface{}, error) { 700 arg := arguments[0].(map[string]interface{}) 701 collected := make([]interface{}, 0, len(arg)) 702 for key := range arg { 703 collected = append(collected, key) 704 } 705 return collected, nil 706 } 707 func jpfValues(arguments []interface{}) (interface{}, error) { 708 arg := arguments[0].(map[string]interface{}) 709 collected := make([]interface{}, 0, len(arg)) 710 for _, value := range arg { 711 collected = append(collected, value) 712 } 713 return collected, nil 714 } 715 func jpfSort(arguments []interface{}) (interface{}, error) { 716 if items, ok := toArrayNum(arguments[0]); ok { 717 d := sort.Float64Slice(items) 718 sort.Stable(d) 719 final := make([]interface{}, len(d)) 720 for i, val := range d { 721 final[i] = val 722 } 723 return final, nil 724 } 725 // Otherwise we're dealing with sort()'ing strings. 726 items, _ := toArrayStr(arguments[0]) 727 d := sort.StringSlice(items) 728 sort.Stable(d) 729 final := make([]interface{}, len(d)) 730 for i, val := range d { 731 final[i] = val 732 } 733 return final, nil 734 } 735 func jpfSortBy(arguments []interface{}) (interface{}, error) { 736 intr := arguments[0].(*treeInterpreter) 737 arr := arguments[1].([]interface{}) 738 exp := arguments[2].(expRef) 739 node := exp.ref 740 if len(arr) == 0 { 741 return arr, nil 742 } else if len(arr) == 1 { 743 return arr, nil 744 } 745 start, err := intr.Execute(node, arr[0]) 746 if err != nil { 747 return nil, err 748 } 749 if _, ok := start.(float64); ok { 750 sortable := &byExprFloat{intr, node, arr, false} 751 sort.Stable(sortable) 752 if sortable.hasError { 753 return nil, errors.New("error in sort_by comparison") 754 } 755 return arr, nil 756 } else if _, ok := start.(string); ok { 757 sortable := &byExprString{intr, node, arr, false} 758 sort.Stable(sortable) 759 if sortable.hasError { 760 return nil, errors.New("error in sort_by comparison") 761 } 762 return arr, nil 763 } else { 764 return nil, errors.New("invalid type, must be number of string") 765 } 766 } 767 func jpfJoin(arguments []interface{}) (interface{}, error) { 768 sep := arguments[0].(string) 769 // We can't just do arguments[1].([]string), we have to 770 // manually convert each item to a string. 771 arrayStr := []string{} 772 for _, item := range arguments[1].([]interface{}) { 773 arrayStr = append(arrayStr, item.(string)) 774 } 775 return strings.Join(arrayStr, sep), nil 776 } 777 func jpfReverse(arguments []interface{}) (interface{}, error) { 778 if s, ok := arguments[0].(string); ok { 779 r := []rune(s) 780 for i, j := 0, len(r)-1; i < len(r)/2; i, j = i+1, j-1 { 781 r[i], r[j] = r[j], r[i] 782 } 783 return string(r), nil 784 } 785 items := arguments[0].([]interface{}) 786 length := len(items) 787 reversed := make([]interface{}, length) 788 for i, item := range items { 789 reversed[length-(i+1)] = item 790 } 791 return reversed, nil 792 } 793 func jpfToArray(arguments []interface{}) (interface{}, error) { 794 if _, ok := arguments[0].([]interface{}); ok { 795 return arguments[0], nil 796 } 797 return arguments[:1:1], nil 798 } 799 func jpfToString(arguments []interface{}) (interface{}, error) { 800 if v, ok := arguments[0].(string); ok { 801 return v, nil 802 } 803 result, err := json.Marshal(arguments[0]) 804 if err != nil { 805 return nil, err 806 } 807 return string(result), nil 808 } 809 func jpfToNumber(arguments []interface{}) (interface{}, error) { 810 arg := arguments[0] 811 if v, ok := arg.(float64); ok { 812 return v, nil 813 } 814 if v, ok := arg.(string); ok { 815 conv, err := strconv.ParseFloat(v, 64) 816 if err != nil { 817 return nil, nil 818 } 819 return conv, nil 820 } 821 if _, ok := arg.([]interface{}); ok { 822 return nil, nil 823 } 824 if _, ok := arg.(map[string]interface{}); ok { 825 return nil, nil 826 } 827 if arg == nil { 828 return nil, nil 829 } 830 if arg == true || arg == false { 831 return nil, nil 832 } 833 return nil, errors.New("unknown type") 834 } 835 func jpfNotNull(arguments []interface{}) (interface{}, error) { 836 for _, arg := range arguments { 837 if arg != nil { 838 return arg, nil 839 } 840 } 841 return nil, nil 842 }