iter_spec.lua (17299B)
1 local t = require('test.testutil') 2 3 local eq = t.eq 4 local matches = t.matches 5 local pcall_err = t.pcall_err 6 7 describe('vim.iter', function() 8 it('new() on iterable class instance', function() 9 local rb = vim.ringbuf(3) 10 rb:push('a') 11 rb:push('b') 12 13 local it = vim.iter(rb) 14 eq({ 'a', 'b' }, it:totable()) 15 end) 16 17 it('filter()', function() 18 local function odd(v) 19 return v % 2 ~= 0 20 end 21 22 local q = { 1, 2, 3, 4, 5 } 23 eq({ 1, 3, 5 }, vim.iter(q):filter(odd):totable()) 24 eq( 25 { 2, 4 }, 26 vim 27 .iter(q) 28 :filter(function(v) 29 return not odd(v) 30 end) 31 :totable() 32 ) 33 eq( 34 {}, 35 vim 36 .iter(q) 37 :filter(function(v) 38 return v > 5 39 end) 40 :totable() 41 ) 42 43 do 44 local it = vim.iter(ipairs(q)) 45 it:filter(function(i, v) 46 return i > 1 and v < 5 47 end) 48 it:map(function(_, v) 49 return v * 2 50 end) 51 eq({ 4, 6, 8 }, it:totable()) 52 end 53 54 local it = vim.iter(string.gmatch('the quick brown fox', '%w+')) 55 eq( 56 { 'the', 'fox' }, 57 it:filter(function(s) 58 return #s <= 3 59 end):totable() 60 ) 61 end) 62 63 it('map()', function() 64 local q = { 1, 2, 3, 4, 5 } 65 eq( 66 { 2, 4, 6, 8, 10 }, 67 vim 68 .iter(q) 69 :map(function(v) 70 return 2 * v 71 end) 72 :totable() 73 ) 74 75 local it = vim.gsplit( 76 [[ 77 Line 1 78 Line 2 79 Line 3 80 Line 4 81 ]], 82 '\n' 83 ) 84 85 eq( 86 { 'Lion 2', 'Lion 4' }, 87 vim 88 .iter(it) 89 :map(function(s) 90 local lnum = s:match('(%d+)') 91 if lnum and tonumber(lnum) % 2 == 0 then 92 return vim.trim(s:gsub('Line', 'Lion')) 93 end 94 end) 95 :totable() 96 ) 97 end) 98 99 it('for loops', function() 100 local q = { 1, 2, 3, 4, 5 } 101 local acc = 0 102 for v in 103 vim.iter(q):map(function(v) 104 return v * 3 105 end) 106 do 107 acc = acc + v 108 end 109 eq(45, acc) 110 end) 111 112 it('totable()', function() 113 do 114 local it = vim.iter({ 1, 2, 3 }):map(function(v) 115 return v, v * v 116 end) 117 eq({ { 1, 1 }, { 2, 4 }, { 3, 9 } }, it:totable()) 118 end 119 120 -- Holes in array-like tables are removed 121 eq({ 1, 2, 3 }, vim.iter({ 1, nil, 2, nil, 3 }):totable()) 122 123 do 124 local it = vim.iter(string.gmatch('1,4,lol,17,blah,2,9,3', '%d+')):map(tonumber) 125 eq({ 1, 4, 17, 2, 9, 3 }, it:totable()) 126 end 127 end) 128 129 it('join()', function() 130 eq('1, 2, 3', vim.iter({ 1, 2, 3 }):join(', ')) 131 eq('a|b|c|d', vim.iter(vim.gsplit('a|b|c|d', '|')):join('|')) 132 end) 133 134 it('next()', function() 135 local it = vim.iter({ 1, 2, 3 }):map(function(v) 136 return 2 * v 137 end) 138 eq(2, it:next()) 139 eq(4, it:next()) 140 eq(6, it:next()) 141 eq(nil, it:next()) 142 end) 143 144 it('rev()', function() 145 eq({ 3, 2, 1 }, vim.iter({ 1, 2, 3 }):rev():totable()) 146 147 local it = vim.iter(string.gmatch('abc', '%w')) 148 matches('rev%(%) requires an array%-like table', pcall_err(it.rev, it)) 149 end) 150 151 it('skip()', function() 152 do 153 local q = { 4, 3, 2, 1 } 154 eq(q, vim.iter(q):skip(0):totable()) 155 eq({ 3, 2, 1 }, vim.iter(q):skip(1):totable()) 156 eq({ 2, 1 }, vim.iter(q):skip(2):totable()) 157 eq({ 1 }, vim.iter(q):skip(#q - 1):totable()) 158 eq({}, vim.iter(q):skip(#q):totable()) 159 eq({}, vim.iter(q):skip(#q + 1):totable()) 160 end 161 162 do 163 local function wrong() 164 return false 165 end 166 167 local function correct() 168 return true 169 end 170 171 local q = { 4, 3, 2, 1 } 172 173 eq({ 4, 3, 2, 1 }, vim.iter(q):skip(wrong):totable()) 174 eq( 175 { 2, 1 }, 176 vim 177 .iter(q) 178 :skip(function(x) 179 return x > 2 180 end) 181 :totable() 182 ) 183 eq({}, vim.iter(q):skip(correct):totable()) 184 end 185 186 do 187 local function skip(n) 188 return vim.iter(vim.gsplit('a|b|c|d', '|')):skip(n):totable() 189 end 190 eq({ 'a', 'b', 'c', 'd' }, skip(0)) 191 eq({ 'b', 'c', 'd' }, skip(1)) 192 eq({ 'c', 'd' }, skip(2)) 193 eq({ 'd' }, skip(3)) 194 eq({}, skip(4)) 195 eq({}, skip(5)) 196 end 197 end) 198 199 it('skip(predicate) preserves first non-matching element', function() 200 local it = vim.iter(vim.gsplit('1|2|3|4', '|')) 201 202 it:skip(function(x) 203 return tonumber(x) < 3 204 end) 205 206 eq('3', it:next()) 207 eq('4', it:next()) 208 end) 209 210 it('skip() followed by peek() works correctly', function() 211 local it = vim.iter(vim.gsplit('a|b|c|d', '|')) 212 213 it:skip(2) 214 215 eq('c', it:peek()) 216 eq('c', it:next()) 217 end) 218 219 it('rskip()', function() 220 do 221 local q = { 4, 3, 2, 1 } 222 eq(q, vim.iter(q):rskip(0):totable()) 223 eq({ 4, 3, 2 }, vim.iter(q):rskip(1):totable()) 224 eq({ 4, 3 }, vim.iter(q):rskip(2):totable()) 225 eq({ 4 }, vim.iter(q):rskip(#q - 1):totable()) 226 eq({}, vim.iter(q):rskip(#q):totable()) 227 eq({}, vim.iter(q):rskip(#q + 1):totable()) 228 end 229 230 local it = vim.iter(vim.gsplit('a|b|c|d', '|')) 231 matches('rskip%(%) requires an array%-like table', pcall_err(it.rskip, it, 0)) 232 end) 233 234 it('slice()', function() 235 local q = { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 } 236 eq({ 3, 4, 5, 6, 7 }, vim.iter(q):slice(3, 7):totable()) 237 eq({}, vim.iter(q):slice(6, 5):totable()) 238 eq({}, vim.iter(q):slice(0, 0):totable()) 239 eq({ 1 }, vim.iter(q):slice(1, 1):totable()) 240 eq({ 1, 2 }, vim.iter(q):slice(1, 2):totable()) 241 eq({ 10 }, vim.iter(q):slice(10, 10):totable()) 242 eq({ 8, 9, 10 }, vim.iter(q):slice(8, 11):totable()) 243 244 local it = vim.iter(vim.gsplit('a|b|c|d', '|')) 245 matches('slice%(%) requires an array%-like table', pcall_err(it.slice, it, 1, 3)) 246 end) 247 248 it('nth()', function() 249 do 250 local q = { 4, 3, 2, 1 } 251 eq(nil, vim.iter(q):nth(0)) 252 eq(4, vim.iter(q):nth(1)) 253 eq(3, vim.iter(q):nth(2)) 254 eq(2, vim.iter(q):nth(3)) 255 eq(1, vim.iter(q):nth(4)) 256 eq(nil, vim.iter(q):nth(5)) 257 end 258 259 do 260 local function nth(n) 261 return vim.iter(vim.gsplit('a|b|c|d', '|')):nth(n) 262 end 263 eq(nil, nth(0)) 264 eq('a', nth(1)) 265 eq('b', nth(2)) 266 eq('c', nth(3)) 267 eq('d', nth(4)) 268 eq(nil, nth(5)) 269 end 270 end) 271 272 it('nth(-x) advances in reverse order starting from end', function() 273 do 274 local q = { 4, 3, 2, 1 } 275 eq(nil, vim.iter(q):nth(0)) 276 eq(1, vim.iter(q):nth(-1)) 277 eq(2, vim.iter(q):nth(-2)) 278 eq(3, vim.iter(q):nth(-3)) 279 eq(4, vim.iter(q):nth(-4)) 280 eq(nil, vim.iter(q):nth(-5)) 281 end 282 283 local it = vim.iter(vim.gsplit('a|b|c|d', '|')) 284 matches('rskip%(%) requires an array%-like table', pcall_err(it.nth, it, -1)) 285 end) 286 287 it('take()', function() 288 local function correct() 289 return true 290 end 291 292 local function wrong() 293 return false 294 end 295 296 do 297 local q = { 4, 3, 2, 1 } 298 eq({}, vim.iter(q):take(0):totable()) 299 eq({ 4 }, vim.iter(q):take(1):totable()) 300 eq({ 4, 3 }, vim.iter(q):take(2):totable()) 301 eq({ 4, 3, 2 }, vim.iter(q):take(3):totable()) 302 eq({ 4, 3, 2, 1 }, vim.iter(q):take(4):totable()) 303 eq({ 4, 3, 2, 1 }, vim.iter(q):take(5):totable()) 304 end 305 306 do 307 local q = { 4, 3, 2, 1 } 308 309 eq({}, vim.iter(q):take(wrong):totable()) 310 eq( 311 { 4, 3 }, 312 vim 313 .iter(q) 314 :take(function(x) 315 return x > 2 316 end) 317 :totable() 318 ) 319 eq({ 4, 3, 2, 1 }, vim.iter(q):take(correct):totable()) 320 end 321 322 do 323 local q = { 4, 3, 2, 1 } 324 eq({ 1, 2, 3 }, vim.iter(q):rev():take(3):totable()) 325 eq({ 2, 3, 4 }, vim.iter(q):take(3):rev():totable()) 326 end 327 328 do 329 local q = { 4, 3, 2, 1 } 330 local it = vim.iter(q) 331 eq({ 4, 3 }, it:take(2):totable()) 332 -- tail is already set from the previous take() 333 eq({ 4, 3 }, it:take(3):totable()) 334 end 335 336 do 337 local it = vim.iter(vim.gsplit('a|b|c|d', '|')) 338 eq({ 'a', 'b' }, it:take(2):totable()) 339 -- non-array iterators are consumed by take() 340 eq({}, it:take(2):totable()) 341 end 342 343 do 344 eq({ 'a', 'b', 'c', 'd' }, vim.iter(vim.gsplit('a|b|c|d', '|')):take(correct):totable()) 345 eq( 346 { 'a', 'b', 'c' }, 347 vim 348 .iter(vim.gsplit('a|b|c|d', '|')) 349 :enumerate() 350 :take(function(i, x) 351 return i < 3 or x == 'c' 352 end) 353 :map(function(_, x) 354 return x 355 end) 356 :totable() 357 ) 358 eq({}, vim.iter(vim.gsplit('a|b|c|d', '|')):take(wrong):totable()) 359 end 360 end) 361 362 it('any()', function() 363 local function odd(v) 364 return v % 2 ~= 0 365 end 366 367 do 368 local q = { 4, 8, 9, 10 } 369 eq(true, vim.iter(q):any(odd)) 370 end 371 372 do 373 local q = { 4, 8, 10 } 374 eq(false, vim.iter(q):any(odd)) 375 end 376 377 do 378 eq( 379 true, 380 vim.iter(vim.gsplit('a|b|c|d', '|')):any(function(s) 381 return s == 'd' 382 end) 383 ) 384 eq( 385 false, 386 vim.iter(vim.gsplit('a|b|c|d', '|')):any(function(s) 387 return s == 'e' 388 end) 389 ) 390 end 391 end) 392 393 it('all()', function() 394 local function odd(v) 395 return v % 2 ~= 0 396 end 397 398 do 399 local q = { 3, 5, 7, 9 } 400 eq(true, vim.iter(q):all(odd)) 401 end 402 403 do 404 local q = { 3, 5, 7, 10 } 405 eq(false, vim.iter(q):all(odd)) 406 end 407 408 do 409 eq( 410 true, 411 vim.iter(vim.gsplit('a|a|a|a', '|')):all(function(s) 412 return s == 'a' 413 end) 414 ) 415 eq( 416 false, 417 vim.iter(vim.gsplit('a|a|a|b', '|')):all(function(s) 418 return s == 'a' 419 end) 420 ) 421 end 422 end) 423 424 it('last()', function() 425 local s = 'abcdefghijklmnopqrstuvwxyz' 426 eq('z', vim.iter(vim.split(s, '')):last()) 427 eq('z', vim.iter(vim.gsplit(s, '')):last()) 428 eq( 429 nil, 430 vim 431 .iter({ 1, 2, 3, 4, 5 }) 432 :filter(function() 433 return false 434 end) 435 :last() 436 ) 437 end) 438 439 it('enumerate()', function() 440 local it = vim.iter(vim.gsplit('abc', '')):enumerate() 441 eq({ 1, 'a' }, { it:next() }) 442 eq({ 2, 'b' }, { it:next() }) 443 eq({ 3, 'c' }, { it:next() }) 444 eq({}, { it:next() }) 445 end) 446 447 it('peek()', function() 448 do 449 local it = vim.iter({ 3, 6, 9, 12 }) 450 eq(3, it:peek()) 451 eq(3, it:peek()) 452 eq(3, it:next()) 453 end 454 455 do 456 local it = vim.iter(vim.gsplit('hi', '')) 457 eq('h', it:peek()) 458 eq('h', it:peek()) 459 eq('h', it:next()) 460 eq('i', it:peek()) 461 eq('i', it:next()) 462 end 463 end) 464 465 it('peek() does not consume on function iterators', function() 466 local it = vim.iter(vim.gsplit('a|b|c', '|')) 467 468 eq('a', it:peek()) 469 eq('a', it:peek()) 470 eq('a', it:next()) 471 eq('b', it:next()) 472 end) 473 474 it('peek() before skip(predicate) does not break iteration', function() 475 local it = vim.iter(vim.gsplit('1|2|3|4', '|')) 476 477 eq('1', it:peek()) 478 479 it:skip(function(x) 480 return tonumber(x) < 3 481 end) 482 483 eq('3', it:next()) 484 end) 485 486 it('multiple peek() calls after next()', function() 487 local it = vim.iter(vim.gsplit('a|b|c', '|')) 488 489 eq('a', it:next()) 490 eq('b', it:peek()) 491 eq('b', it:peek()) 492 eq('b', it:next()) 493 eq('c', it:next()) 494 end) 495 496 describe('peek() with multi-value returns', function() 497 it('peek() preserves multiple return values from ipairs()', function() 498 local it = vim.iter(ipairs({ 'a', 'b', 'c' })) 499 local i1, v1 = it:peek() 500 501 eq(1, i1) 502 eq('a', v1) 503 504 local i2, v2 = it:next() 505 506 eq(1, i2) 507 eq('a', v2) 508 end) 509 510 it('peek() works with pairs() returning multiple values', function() 511 local tbl = { x = 10, y = 20 } 512 local it = vim.iter(pairs(tbl)) 513 local k1, v1 = it:peek() 514 local k2, v2 = it:peek() 515 516 eq(k1, k2) 517 eq(v1, v2) 518 end) 519 end) 520 521 describe('peek() after transformations', function() 522 it('peek() works after map() on function iterator', function() 523 local it = vim.iter(vim.gsplit('1|2|3', '|')):map(tonumber) 524 525 eq(1, it:peek()) 526 eq(1, it:next()) 527 eq(2, it:peek()) 528 end) 529 530 it('peek() at end of iterator returns nil', function() 531 local it = vim.iter({ 1 }) 532 533 eq(1, it:next()) 534 eq(nil, it:peek()) 535 eq(nil, it:next()) 536 end) 537 end) 538 539 it('find()', function() 540 local q = { 3, 6, 9, 12 } 541 eq(12, vim.iter(q):find(12)) 542 eq(nil, vim.iter(q):find(15)) 543 eq( 544 12, 545 vim.iter(q):find(function(v) 546 return v % 4 == 0 547 end) 548 ) 549 550 do 551 local it = vim.iter(q) 552 local pred = function(v) 553 return v % 3 == 0 554 end 555 eq(3, it:find(pred)) 556 eq(6, it:find(pred)) 557 eq(9, it:find(pred)) 558 eq(12, it:find(pred)) 559 eq(nil, it:find(pred)) 560 end 561 562 do 563 local it = vim.iter(vim.gsplit('AbCdE', '')) 564 local pred = function(s) 565 return s:match('[A-Z]') 566 end 567 eq('A', it:find(pred)) 568 eq('C', it:find(pred)) 569 eq('E', it:find(pred)) 570 eq(nil, it:find(pred)) 571 end 572 end) 573 574 it('rfind()', function() 575 local q = { 1, 2, 3, 2, 1 } 576 do 577 local it = vim.iter(q) 578 eq(1, it:rfind(1)) 579 eq(1, it:rfind(1)) 580 eq(nil, it:rfind(1)) 581 end 582 583 do 584 local it = vim.iter(q):enumerate() 585 local pred = function(i) 586 return i % 2 ~= 0 587 end 588 eq({ 5, 1 }, { it:rfind(pred) }) 589 eq({ 3, 3 }, { it:rfind(pred) }) 590 eq({ 1, 1 }, { it:rfind(pred) }) 591 eq(nil, it:rfind(pred)) 592 end 593 594 do 595 local it = vim.iter(vim.gsplit('AbCdE', '')) 596 matches('rfind%(%) requires an array%-like table', pcall_err(it.rfind, it, 'E')) 597 end 598 end) 599 600 it('pop()', function() 601 do 602 local it = vim.iter({ 1, 2, 3, 4 }) 603 eq(4, it:pop()) 604 eq(3, it:pop()) 605 eq(2, it:pop()) 606 eq(1, it:pop()) 607 eq(nil, it:pop()) 608 eq(nil, it:pop()) 609 end 610 611 do 612 local it = vim.iter(vim.gsplit('hi', '')) 613 matches('pop%(%) requires an array%-like table', pcall_err(it.pop, it)) 614 end 615 end) 616 617 it('rpeek()', function() 618 do 619 local it = vim.iter({ 1, 2, 3, 4 }) 620 eq(4, it:rpeek()) 621 eq(4, it:rpeek()) 622 eq(4, it:pop()) 623 end 624 625 do 626 local it = vim.iter(vim.gsplit('hi', '')) 627 matches('rpeek%(%) requires an array%-like table', pcall_err(it.rpeek, it)) 628 end 629 end) 630 631 it('fold()', function() 632 local q = { 1, 2, 3, 4, 5 } 633 eq( 634 115, 635 vim.iter(q):fold(100, function(acc, v) 636 return acc + v 637 end) 638 ) 639 eq( 640 { 5, 4, 3, 2, 1 }, 641 vim.iter(q):fold({}, function(acc, v) 642 table.insert(acc, 1, v) 643 return acc 644 end) 645 ) 646 end) 647 648 it('flatten()', function() 649 local q = { { 1, { 2 } }, { { { { 3 } } }, { 4 } }, { 5 } } 650 651 eq(q, vim.iter(q):flatten(-1):totable()) 652 eq(q, vim.iter(q):flatten(0):totable()) 653 eq({ 1, { 2 }, { { { 3 } } }, { 4 }, 5 }, vim.iter(q):flatten():totable()) 654 eq({ 1, 2, { { 3 } }, 4, 5 }, vim.iter(q):flatten(2):totable()) 655 eq({ 1, 2, { 3 }, 4, 5 }, vim.iter(q):flatten(3):totable()) 656 eq({ 1, 2, 3, 4, 5 }, vim.iter(q):flatten(4):totable()) 657 658 local m = { a = 1, b = { 2, 3 }, d = { 4 } } 659 local it = vim.iter(m) 660 661 local flat_err = 'flatten%(%) requires an array%-like table' 662 matches(flat_err, pcall_err(it.flatten, it)) 663 664 -- cases from the documentation 665 local simple_example = { 1, { 2 }, { { 3 } } } 666 eq({ 1, 2, { 3 } }, vim.iter(simple_example):flatten():totable()) 667 668 local not_list_like = { [2] = 2 } 669 eq({ 2 }, vim.iter(not_list_like):flatten():totable()) 670 671 local also_not_list_like = { nil, 2 } 672 eq({ 2 }, vim.iter(also_not_list_like):flatten():totable()) 673 674 eq({ 1, 2, 3 }, vim.iter({ nil, { 1, nil, 2 }, 3 }):flatten():totable()) 675 676 local nested_non_lists = vim.iter({ 1, { { a = 2 } }, { { nil } }, { 3 } }) 677 eq({ 1, { a = 2 }, { nil }, 3 }, nested_non_lists:flatten():totable()) 678 -- only error if we're going deep enough to flatten a dict-like table 679 matches(flat_err, pcall_err(nested_non_lists.flatten, nested_non_lists, math.huge)) 680 end) 681 682 it('unique()', function() 683 eq({ 1, 2, 3, 4, 5 }, vim.iter({ 1, 2, 2, 3, 4, 4, 5 }):unique():totable()) 684 eq( 685 { 1, 2, 3, 4, 5 }, 686 vim.iter({ 1, 2, 3, 4, 4, 5, 1, 2, 3, 2, 1, 2, 3, 4, 5 }):unique():totable() 687 ) 688 eq( 689 { { 1 }, { 2 }, { 3 } }, 690 vim 691 .iter({ { 1 }, { 1 }, { 2 }, { 2 }, { 3 }, { 3 } }) 692 :unique(function(x) 693 return x[1] 694 end) 695 :totable() 696 ) 697 end) 698 699 it('handles map-like tables', function() 700 local it = vim.iter({ a = 1, b = 2, c = 3 }):map(function(k, v) 701 if v % 2 ~= 0 then 702 return k:upper(), v * 2 703 end 704 end) 705 706 local q = it:fold({}, function(q, k, v) 707 q[k] = v 708 return q 709 end) 710 eq({ A = 2, C = 6 }, q) 711 end) 712 713 it('handles table values mid-pipeline', function() 714 local map = { 715 item = { 716 file = 'test', 717 }, 718 item_2 = { 719 file = 'test', 720 }, 721 item_3 = { 722 file = 'test', 723 }, 724 } 725 726 local output = vim 727 .iter(map) 728 :map(function(key, value) 729 return { [key] = value.file } 730 end) 731 :totable() 732 733 table.sort(output, function(a, b) 734 return next(a) < next(b) 735 end) 736 737 eq({ 738 { item = 'test' }, 739 { item_2 = 'test' }, 740 { item_3 = 'test' }, 741 }, output) 742 end) 743 end)