Skip to content

Commit 7c63e02

Browse files
add tets
Signed-off-by: OneSizeFitsQuorum <[email protected]>
1 parent 2b14563 commit 7c63e02

File tree

2 files changed

+49
-2
lines changed

2 files changed

+49
-2
lines changed

fsspec/callbacks.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,25 @@ def set_size(self, size):
9292
Parameters
9393
----------
9494
size: int or callable
95+
The total size of the transfer. Can be either:
96+
- An integer representing the total size directly
97+
- A callable (function/method) that returns an integer when invoked
98+
99+
The callable option is useful when the size is only available as a
100+
method on an object (e.g., filesystem objects that have a ``size()``
101+
method instead of a ``size`` attribute).
102+
103+
Examples
104+
--------
105+
>>> callback = Callback()
106+
>>> callback.set_size(1000) # Direct integer
107+
>>> callback.set_size(lambda: 1000) # Callable returning integer
108+
109+
Notes
110+
-----
111+
If a callable is provided, it will be invoked immediately to obtain
112+
the size value. The callable should take no arguments and return an
113+
integer.
95114
"""
96115
if callable(size):
97116
size = size()
@@ -323,4 +342,4 @@ def __del__(self):
323342
return self.close()
324343

325344

326-
DEFAULT_CALLBACK = _DEFAULT_CALLBACK = NoOpCallback()
345+
DEFAULT_CALLBACK = _DEFAULT_CALLBACK = NoOpCallback()

fsspec/tests/test_callbacks.py

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,34 @@ def relative_update(self, inc=1):
7373
assert events == [1] * 10
7474

7575

76+
def test_set_size_with_callable():
77+
"""Test that set_size accepts both int and callable parameters."""
78+
callback = Callback()
79+
80+
# Test with integer
81+
callback.set_size(100)
82+
assert callback.size == 100
83+
84+
# Test with callable (lambda)
85+
callback.set_size(lambda: 200)
86+
assert callback.size == 200
87+
88+
# Test with callable (function)
89+
def get_size():
90+
return 300
91+
callback.set_size(get_size)
92+
assert callback.size == 300
93+
94+
# Test with callable that simulates a method attribute
95+
class MockFileSystem:
96+
def size(self):
97+
return 400
98+
99+
fs = MockFileSystem()
100+
callback.set_size(fs.size)
101+
assert callback.size == 400
102+
103+
76104
@pytest.mark.parametrize("tqdm_kwargs", [{}, {"desc": "A custom desc"}])
77105
def test_tqdm_callback(tqdm_kwargs, mocker):
78106
pytest.importorskip("tqdm")
@@ -86,4 +114,4 @@ def test_tqdm_callback(tqdm_kwargs, mocker):
86114
if not tqdm_kwargs:
87115
callback._tqdm_cls.assert_called_with(total=10)
88116
else:
89-
callback._tqdm_cls.assert_called_with(total=10, **tqdm_kwargs)
117+
callback._tqdm_cls.assert_called_with(total=10, **tqdm_kwargs)

0 commit comments

Comments
 (0)