Skip to content

Commit 5f8c5d6

Browse files
feat: Support non __len__ and non __getitem__
1 parent 21165c4 commit 5f8c5d6

1 file changed

Lines changed: 27 additions & 0 deletions

File tree

src/thread/thread.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)