@@ -357,6 +357,7 @@ class ParallelProcessing(Generic[_Target_P, _Target_T, _Dataset_T]):
357357 SupportsLengthGetItem [_Dataset_T ],
358358 SupportsGetItem [_Dataset_T ],
359359 SupportsLength ,
360+ type [_Dataset_T ],
360361 ]
361362 max_threads : int
362363
@@ -408,13 +409,27 @@ def __init__(
408409 ** overflow_kwargs : Overflow_In ,
409410 ) -> None : ...
410411
412+ # Does not support __getitem__ and __len__
413+ @overload
414+ def __init__ (
415+ self ,
416+ function : DatasetFunction [_Dataset_T , _Target_P , _Target_T ],
417+ dataset : type [_Dataset_T ],
418+ max_threads : int = 8 ,
419+ * overflow_args : Overflow_In ,
420+ _get_value : Callable [[type [_Dataset_T ], int ], _Dataset_T ],
421+ _length : Union [int , Callable [[type [_Dataset_T ]], int ]],
422+ ** overflow_kwargs : Overflow_In ,
423+ ) -> None : ...
424+
411425 def __init__ (
412426 self ,
413427 function : DatasetFunction [_Dataset_T , _Target_P , _Target_T ],
414428 dataset : Union [
415429 SupportsLengthGetItem [_Dataset_T ],
416430 SupportsGetItem [_Dataset_T ],
417431 SupportsLength ,
432+ type [_Dataset_T ],
418433 ],
419434 max_threads : int = 8 ,
420435 * overflow_args : Overflow_In ,
@@ -423,6 +438,7 @@ def __init__(
423438 Callable [[SupportsLengthGetItem [_Dataset_T ], int ], _Dataset_T ],
424439 Callable [[SupportsGetItem [_Dataset_T ], int ], _Dataset_T ],
425440 Callable [[SupportsLength , int ], _Dataset_T ],
441+ Callable [[type [_Dataset_T ], int ], _Dataset_T ],
426442 ]
427443 ] = None ,
428444 _length : Optional [
@@ -480,6 +496,17 @@ def __init__(
480496
481497 get_value = _get_value or dataset .__class__ .__getitem__
482498
499+ else :
500+ assert (
501+ _length
502+ ), '`_length` must be set if `dataset` does not support `__len__`'
503+ assert (
504+ _get_value
505+ ), '`_get_value` must be set if `dataset` does not support `__getitem__`'
506+
507+ length = _length
508+ get_value = _get_value
509+
483510 assert isinstance (length , int ), '`_length` must be an integer'
484511 assert length > 0 , 'dataset cannot be empty'
485512 assert get_value , '`_get_value` must be set'
0 commit comments