From 66d37e0be2a17e3fb6812c12ddeef5ce4676f452 Mon Sep 17 00:00:00 2001
From: simon987 <me@simon987.net>
Date: Sun, 28 Feb 2021 12:15:44 -0500
Subject: [PATCH] Add way to manually flush @buffered

---
 hexlib/misc.py        |  6 ++++++
 setup.py              |  2 +-
 test/test_buffered.py | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 test/test_buffered.py

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)