@@ -12,12 +12,17 @@ local f = require "plenary.functional"
1212-- Tools
1313---- ----------------------------------------------------------------------------
1414
15+ --- @class PlenaryIterators
1516local exports = {}
1617
17- --- @class Iterator
18- --- @field gen function
19- --- @field param any
20- --- @field state any
18+ --- @generic V
19+ --- @alias PlenaryIteratorsIterator fun ( table : V[] , i ?: integer ): integer , V ?
20+
21+ --- @class PlenaryIterator
22+ --- @field gen PlenaryIteratorsIterator
23+ --- @field param table
24+ --- @field state ? integer
25+ --- @overload fun ( param ?: table , state ?: integer ): integer , any ?
2126local Iterator = {}
2227Iterator .__index = Iterator
2328
@@ -34,6 +39,7 @@ function Iterator:__call(param, state)
3439 return self .gen (param or self .param , state or self .state )
3540end
3641
42+ --- @return string
3743function Iterator :__tostring ()
3844 return " <iterator>"
3945end
@@ -86,6 +92,10 @@ local map_gen = function(map, key)
8692 return key , key , value
8793end
8894
95+ --- @param param string
96+ --- @param state integer
97+ --- @return integer ? state
98+ --- @return string ? r
8999local string_gen = function (param , state )
90100 state = state + 1
91101 if state > # param then
@@ -95,6 +105,18 @@ local string_gen = function(param, state)
95105 return state , r
96106end
97107
108+ --- @generic T : table , U : table
109+ --- @alias PlenaryIteratorsRawiterTable fun ( obj : T | PlenaryIterator , param ?: string , state ?: integer ): PlenaryIteratorsIterator , T | U | nil , integer ?
110+
111+ --- @generic T : table , U : table
112+ --- @param obj T | PlenaryIterator
113+ --- @param param ? string
114+ --- @param state ? integer
115+ --- @return PlenaryIteratorsIterator gen
116+ --- @return T | U | nil param
117+ --- @return integer ? state
118+ --- @overload fun ( obj : PlenaryIteratorsIterator , param : any , state ?: integer ): PlenaryIteratorsIterator , any , integer ?
119+ --- @overload fun ( obj : string ): PlenaryIteratorsIterator , string ?, integer ?
98120local rawiter = function (obj , param , state )
99121 assert (obj ~= nil , " invalid iterator" )
100122
132154--- @param gen any
133155--- @param param any
134156--- @param state any
135- --- @return Iterator
157+ --- @return PlenaryIterator
136158local function wrap (gen , param , state )
137159 return setmetatable ({
138160 gen = gen ,
@@ -142,7 +164,7 @@ local function wrap(gen, param, state)
142164end
143165
144166--- Unwrap an iterator metatable into the iterator triplet
145- --- @param self Iterator
167+ --- @param self PlenaryIterator
146168--- @return any
147169--- @return any
148170--- @return any
154176--- @param obj any
155177--- @param param any (optional )
156178--- @param state any (optional )
157- --- @return Iterator
179+ --- @return PlenaryIterator
158180local iter = function (obj , param , state )
159181 return wrap (rawiter (obj , param , state ))
160182end
228250--- @param start number
229251--- @param stop number
230252--- @param step number
231- --- @return Iterator
253+ --- @return PlenaryIterator
232254local range = function (start , stop , step )
233255 if step == nil then
234256 if stop == nil then
269291--- Creates an infinite iterator that will yield the arguments
270292--- If multiple arguments are passed, the args will be packed and unpacked
271293--- @param ...: the arguments to duplicate
272- --- @return Iterator
294+ --- @return PlenaryIterator
273295local duplicate = function (...)
274296 if select (" #" , ... ) <= 1 then
275297 return wrap (duplicate_gen , select (1 , ... ), 0 )
@@ -282,7 +304,7 @@ exports.duplicate = duplicate
282304--- Creates an iterator from a function
283305--- NOTE: if the function is a closure and modifies state, the resulting iterator will not be stateless
284306--- @param fun function
285- --- @return Iterator
307+ --- @return PlenaryIterator
286308local from_fun = function (fun )
287309 assert (type (fun ) == " function" )
288310 return wrap (duplicate_fun_gen , fun , 0 )
@@ -291,15 +313,15 @@ exports.from_fun = from_fun
291313
292314--- Creates an infinite iterator that will yield zeros.
293315--- This is an alias to calling duplicate(0)
294- --- @return Iterator
316+ --- @return PlenaryIterator
295317local zeros = function ()
296318 return wrap (duplicate_gen , 0 , 0 )
297319end
298320exports .zeros = zeros
299321
300322--- Creates an infinite iterator that will yield ones.
301323--- This is an alias to calling duplicate(1)
302- --- @return Iterator
324+ --- @return PlenaryIterator
303325local ones = function ()
304326 return wrap (duplicate_gen , 1 , 0 )
305327end
316338--- Creates an infinite iterator that will yield random values.
317339--- @param n number
318340--- @param m number
319- --- @return Iterator
341+ --- @return PlenaryIterator
320342local rands = function (n , m )
321343 if n == nil and m == nil then
322344 return wrap (rands_nil_gen , 0 , 0 )
355377--- Return an iterator of substrings separated by a string
356378--- @param input string : the string to split
357379--- @param sep string : the separator to find and split based on
358- --- @return Iterator
380+ --- @return PlenaryIterator
359381local split = function (input , sep )
360382 return wrap (split_gen , { input , sep }, 1 )
361383end
386408
387409--- Iterator adapter that maps the previous iterator with a function
388410--- @param fun function : The function to map with. Will be called on each element
389- --- @return Iterator
411+ --- @return PlenaryIterator
390412function Iterator :map (fun )
391413 return wrap (map_gen2 , { self .gen , self .param , fun }, self .state )
392414end
@@ -428,7 +450,7 @@ local flatten_gen = function(_, state)
428450end
429451
430452--- Iterator adapter that will recursivley flatten nested iterator structure
431- --- @return Iterator
453+ --- @return PlenaryIterator
432454function Iterator :flatten ()
433455 return wrap (flatten_gen , false , { self .gen , self .param , self .state })
434456end
@@ -481,13 +503,13 @@ end
481503
482504--- Iterator adapter that will filter values
483505--- @param fun function : The function to filter values with. If the function returns true , the value will be kept.
484- --- @return Iterator
506+ --- @return PlenaryIterator
485507function Iterator :filter (fun )
486508 return wrap (filter_gen , { self .gen , self .param , fun }, self .state )
487509end
488510
489511--- Iterator adapter that will provide numbers from 1 to n as the first multival
490- --- @return Iterator
512+ --- @return PlenaryIterator
491513function Iterator :enumerate ()
492514 local i = 0
493515 return self :map (function (...)
617639--- Used for treating consecutive iterators as a single iterator.
618640--- Infinity iterators are supported, but are not recommended.
619641--- @param ...: the iterators to chain
620- --- @return Iterator
642+ --- @return PlenaryIterator
621643local chain = function (...)
622644 local n = numargs (... )
623645
666688--- The returned iterator is truncated in length to the length of the shortest iterator.
667689--- For multi-return iterators only the first variable is used.
668690--- @param ...: the iterators to zip
669- --- @return Iterator
691+ --- @return PlenaryIterator
670692local zip = function (...)
671693 local n = numargs (... )
672694 if n == 0 then
0 commit comments