diff --git a/hexlib/misc.py b/hexlib/misc.py index fd7981a..127aaf7 100644 --- a/hexlib/misc.py +++ b/hexlib/misc.py @@ -57,6 +57,12 @@ def buffered(batch_size: int, flush_on_exit: bool = False): atexit.register(func, buffer) def wrapper(items): + + if items is None: + func(buffer) + buffer.clear() + return + with lock: for item in items: buffer.append(item) diff --git a/setup.py b/setup.py index 0f6dc50..3f1a104 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ from setuptools import setup setup( name="hexlib", - version="1.29", + version="1.30", description="Misc utility methods", author="simon987", author_email="me@simon987.net", diff --git a/test/test_buffered.py b/test/test_buffered.py new file mode 100644 index 0000000..a42047f --- /dev/null +++ b/test/test_buffered.py @@ -0,0 +1,36 @@ +from unittest import TestCase + +from hexlib.misc import buffered + + +class TestBuffered(TestCase): + + def test_simple(self): + my_list = [] + + @buffered(batch_size=2) + def put_item(items): + my_list.extend(items) + + put_item([1, 2]) + put_item([1]) + put_item([1]) + put_item([1]) + + self.assertEqual(len(my_list), 4) + + def test_flush(self): + my_list = [] + + @buffered(batch_size=2) + def put_item(items): + my_list.extend(items) + + put_item([1, 2]) + put_item([1]) + put_item([1]) + put_item([1]) + + put_item(None) + + self.assertEqual(len(my_list), 5)