@@ -13,12 +13,17 @@ local compat = require "plenary.compat"
1313-- Tools
1414---- ----------------------------------------------------------------------------
1515
16+ --- @class PlenaryIterators
1617local exports = {}
1718
18- --- @class Iterator
19- --- @field gen function
20- --- @field param any
21- --- @field state any
19+ --- @generic V
20+ --- @alias PlenaryIteratorsIterator fun ( table : V[] , i ?: integer ): integer , V ?
21+
22+ --- @class PlenaryIterator
23+ --- @field gen PlenaryIteratorsIterator
24+ --- @field param table
25+ --- @field state ? integer
26+ --- @overload fun ( param ?: table , state ?: integer ): integer , any ?
2227local Iterator = {}
2328Iterator .__index = Iterator
2429
@@ -35,6 +40,7 @@ function Iterator:__call(param, state)
3540 return self .gen (param or self .param , state or self .state )
3641end
3742
43+ --- @return string
3844function Iterator :__tostring ()
3945 return " <iterator>"
4046end
@@ -87,6 +93,10 @@ local map_gen = function(map, key)
8793 return key , key , value
8894end
8995
96+ --- @param param string
97+ --- @param state integer
98+ --- @return integer ? state
99+ --- @return string ? r
90100local string_gen = function (param , state )
91101 state = state + 1
92102 if state > # param then
@@ -96,6 +106,18 @@ local string_gen = function(param, state)
96106 return state , r
97107end
98108
109+ --- @generic T : table , U : table
110+ --- @alias PlenaryIteratorsRawiterTable fun ( obj : T | PlenaryIterator , param ?: string , state ?: integer ): PlenaryIteratorsIterator , T | U | nil , integer ?
111+
112+ --- @generic T : table , U : table
113+ --- @param obj T | PlenaryIterator
114+ --- @param param ? string
115+ --- @param state ? integer
116+ --- @return PlenaryIteratorsIterator gen
117+ --- @return T | U | nil param
118+ --- @return integer ? state
119+ --- @overload fun ( obj : PlenaryIteratorsIterator , param : any , state ?: integer ): PlenaryIteratorsIterator , any , integer ?
120+ --- @overload fun ( obj : string ): PlenaryIteratorsIterator , string ?, integer ?
99121local rawiter = function (obj , param , state )
100122 assert (obj ~= nil , " invalid iterator" )
101123
133155--- @param gen any
134156--- @param param any
135157--- @param state any
136- --- @return Iterator
158+ --- @return PlenaryIterator
137159local function wrap (gen , param , state )
138160 return setmetatable ({
139161 gen = gen ,
@@ -143,7 +165,7 @@ local function wrap(gen, param, state)
143165end
144166
145167--- Unwrap an iterator metatable into the iterator triplet
146- --- @param self Iterator
168+ --- @param self PlenaryIterator
147169--- @return any
148170--- @return any
149171--- @return any
155177--- @param obj any
156178--- @param param any (optional )
157179--- @param state any (optional )
158- --- @return Iterator
180+ --- @return PlenaryIterator
159181local iter = function (obj , param , state )
160182 return wrap (rawiter (obj , param , state ))
161183end
229251--- @param start number
230252--- @param stop number
231253--- @param step number
232- --- @return Iterator
254+ --- @return PlenaryIterator
233255local range = function (start , stop , step )
234256 if step == nil then
235257 if stop == nil then
270292--- Creates an infinite iterator that will yield the arguments
271293--- If multiple arguments are passed, the args will be packed and unpacked
272294--- @param ...: the arguments to duplicate
273- --- @return Iterator
295+ --- @return PlenaryIterator
274296local duplicate = function (...)
275297 if select (" #" , ... ) <= 1 then
276298 return wrap (duplicate_gen , select (1 , ... ), 0 )
@@ -283,7 +305,7 @@ exports.duplicate = duplicate
283305--- Creates an iterator from a function
284306--- NOTE: if the function is a closure and modifies state, the resulting iterator will not be stateless
285307--- @param fun function
286- --- @return Iterator
308+ --- @return PlenaryIterator
287309local from_fun = function (fun )
288310 assert (type (fun ) == " function" )
289311 return wrap (duplicate_fun_gen , fun , 0 )
@@ -292,15 +314,15 @@ exports.from_fun = from_fun
292314
293315--- Creates an infinite iterator that will yield zeros.
294316--- This is an alias to calling duplicate(0)
295- --- @return Iterator
317+ --- @return PlenaryIterator
296318local zeros = function ()
297319 return wrap (duplicate_gen , 0 , 0 )
298320end
299321exports .zeros = zeros
300322
301323--- Creates an infinite iterator that will yield ones.
302324--- This is an alias to calling duplicate(1)
303- --- @return Iterator
325+ --- @return PlenaryIterator
304326local ones = function ()
305327 return wrap (duplicate_gen , 1 , 0 )
306328end
317339--- Creates an infinite iterator that will yield random values.
318340--- @param n number
319341--- @param m number
320- --- @return Iterator
342+ --- @return PlenaryIterator
321343local rands = function (n , m )
322344 if n == nil and m == nil then
323345 return wrap (rands_nil_gen , 0 , 0 )
356378--- Return an iterator of substrings separated by a string
357379--- @param input string : the string to split
358380--- @param sep string : the separator to find and split based on
359- --- @return Iterator
381+ --- @return PlenaryIterator
360382local split = function (input , sep )
361383 return wrap (split_gen , { input , sep }, 1 )
362384end
387409
388410--- Iterator adapter that maps the previous iterator with a function
389411--- @param fun function : The function to map with. Will be called on each element
390- --- @return Iterator
412+ --- @return PlenaryIterator
391413function Iterator :map (fun )
392414 return wrap (map_gen2 , { self .gen , self .param , fun }, self .state )
393415end
@@ -429,7 +451,7 @@ local flatten_gen = function(_, state)
429451end
430452
431453--- Iterator adapter that will recursivley flatten nested iterator structure
432- --- @return Iterator
454+ --- @return PlenaryIterator
433455function Iterator :flatten ()
434456 return wrap (flatten_gen , false , { self .gen , self .param , self .state })
435457end
@@ -482,13 +504,13 @@ end
482504
483505--- Iterator adapter that will filter values
484506--- @param fun function : The function to filter values with. If the function returns true , the value will be kept.
485- --- @return Iterator
507+ --- @return PlenaryIterator
486508function Iterator :filter (fun )
487509 return wrap (filter_gen , { self .gen , self .param , fun }, self .state )
488510end
489511
490512--- Iterator adapter that will provide numbers from 1 to n as the first multival
491- --- @return Iterator
513+ --- @return PlenaryIterator
492514function Iterator :enumerate ()
493515 local i = 0
494516 return self :map (function (...)
618640--- Used for treating consecutive iterators as a single iterator.
619641--- Infinity iterators are supported, but are not recommended.
620642--- @param ...: the iterators to chain
621- --- @return Iterator
643+ --- @return PlenaryIterator
622644local chain = function (...)
623645 local n = numargs (... )
624646
667689--- The returned iterator is truncated in length to the length of the shortest iterator.
668690--- For multi-return iterators only the first variable is used.
669691--- @param ...: the iterators to zip
670- --- @return Iterator
692+ --- @return PlenaryIterator
671693local zip = function (...)
672694 local n = numargs (... )
673695 if n == 0 then
0 commit comments