Skip to content

Commit ce70b45

Browse files
committed
Migrate to doctest
1 parent 70a0f3c commit ce70b45

File tree

1 file changed

+29
-16
lines changed

1 file changed

+29
-16
lines changed

torchdata/datapipes/iter/util/converter.py

Lines changed: 29 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -36,24 +36,37 @@ class IterToMapConverterMapDataPipe(MapDataPipe):
3636
will be replaced by the new value.
3737
3838
Example:
39-
>>> from torchdata.datapipes.iter import IterableWrapper
40-
>>> source_dp = IterableWrapper([(i, i) for i in range(10)])
41-
>>> map_dp = source_dp.to_map_datapipe()
42-
>>> list(map_dp)
43-
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
44-
>>> source_dp2 = IterableWrapper([('a', 1), ('b', 2), ('c', 1)])
45-
>>> map_dp2 = source_dp2.to_map_datapipe()
46-
>>> map_dp2['a']
47-
1
48-
>>> def row_to_tuple(row):
49-
>>> label = row[0]
50-
>>> data = row[1:]
51-
>>> return label, data
52-
>>> source_dp3 = IterableWrapper([('a', 1, 1, 1, 1, 1, 1), ('b', 2, 2, 2, 2, 2, 2), ('c', 3, 3, 3, 3, 3, 3)])
53-
>>> map_dp3 = source_dp3.to_map_datapipe(key_value_fn=row_to_tuple)
54-
>>> map_dp3['a']
39+
40+
.. testsetup::
41+
42+
from torchdata.datapipes.iter import IterableWrapper
43+
44+
.. testcode::
45+
46+
source_dp = IterableWrapper([(i, i) for i in range(10)])
47+
map_dp = source_dp.to_map_datapipe()
48+
assert list(map_dp) == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
49+
50+
source_dp2 = IterableWrapper([('a', 1), ('b', 2), ('c', 1)])
51+
map_dp2 = source_dp2.to_map_datapipe()
52+
assert map_dp2['a']) == 1
53+
54+
.. testcode::
55+
56+
def row_to_tuple(row):
57+
label = row[0]
58+
data = row[1:]
59+
return label, data
60+
source_dp3 = IterableWrapper([('a', 1, 1, 1, 1, 1, 1), ('b', 2, 2, 2, 2, 2, 2), ('c', 3, 3, 3, 3, 3, 3)])
61+
map_dp3 = source_dp3.to_map_datapipe(key_value_fn=row_to_tuple)
62+
print(map_dp3['a'])
63+
64+
.. testoutput::
65+
5566
(1, 1, 1, 1, 1, 1)
67+
5668
"""
69+
5770
datapipe: IterDataPipe
5871
key_value_fn: Optional[Callable]
5972
_map: Optional[Dict]

0 commit comments

Comments
 (0)