add test

Signed-off-by: OneSizeFitsQuorum <[email protected]>

Author: OneSizeFitsQuorum

fsspec/callbacks.py

@@ -92,6 +92,25 @@ def set_size(self, size):
         Parameters
         ----------
         size: int or callable
+            The total size of the transfer. Can be either:
+            - An integer representing the total size directly
+            - A callable (function/method) that returns an integer when invoked
+
+            The callable option is useful when the size is only available as a
+            method on an object (e.g., filesystem objects that have a ``size()``
+            method instead of a ``size`` attribute).
+
+        Examples
+        --------
+        >>> callback = Callback()
+        >>> callback.set_size(1000)  # Direct integer
+        >>> callback.set_size(lambda: 1000)  # Callable returning integer
+
+        Notes
+        -----
+        If a callable is provided, it will be invoked immediately to obtain
+        the size value. The callable should take no arguments and return an
+        integer.
         """
         if callable(size):
             size = size()

fsspec/tests/test_callbacks.py

@@ -73,6 +73,35 @@ def relative_update(self, inc=1):
     assert events == [1] * 10
 
 
+def test_set_size_with_callable():
+    """Test that set_size accepts both int and callable parameters."""
+    callback = Callback()
+
+    # Test with integer
+    callback.set_size(100)
+    assert callback.size == 100
+
+    # Test with callable (lambda)
+    callback.set_size(lambda: 200)
+    assert callback.size == 200
+
+    # Test with callable (function)
+    def get_size():
+        return 300
+
+    callback.set_size(get_size)
+    assert callback.size == 300
+
+    # Test with callable that simulates a method attribute
+    class MockFileSystem:
+        def size(self):
+            return 400
+
+    fs = MockFileSystem()
+    callback.set_size(fs.size)
+    assert callback.size == 400
+
+
 @pytest.mark.parametrize("tqdm_kwargs", [{}, {"desc": "A custom desc"}])
 def test_tqdm_callback(tqdm_kwargs, mocker):
     pytest.importorskip("tqdm")