From 52ee1dab95436128b44c37cc495022ff90108b2e Mon Sep 17 00:00:00 2001
From: Nikolaus Rath <Nikolaus@rath.org>
Date: Mon, 9 May 2022 10:04:30 +0100
Subject: [PATCH] Add nbdkit.parse_size() Python function.

This enables Python plugins to parse sizes the same way as C plugins.

I'm not sure about the best way to test this - input is appreciated.

I'm not too happy with the way this code is tested. It workes, but putting the tests into
test-python-plugin.py feels misplaced: this file is intended to support the unit tests in
test_python.py, not run its own unit tests.

(cherry picked from commit 1b7d72542be68e254c1ef86ecb1a82b05c78ff63)
---
 plugins/python/modfunctions.c           | 21 +++++++++++++++++++++
 plugins/python/nbdkit-python-plugin.pod |  5 +++++
 tests/test-python-plugin.py             | 19 +++++++++++++++++++
 3 files changed, 45 insertions(+)

diff --git a/plugins/python/modfunctions.c b/plugins/python/modfunctions.c
index fffbaab2..46b0c904 100644
--- a/plugins/python/modfunctions.c
+++ b/plugins/python/modfunctions.c
@@ -93,11 +93,32 @@ do_shutdown (PyObject *self, PyObject *args)
   Py_RETURN_NONE;
 }
 
+/* nbdkit.parse_size */
+static PyObject *
+parse_size (PyObject *self, PyObject *args)
+{
+  const char *s;
+  if (!PyArg_ParseTuple (args, "s", &s)) {
+    PyErr_SetString (PyExc_TypeError, "Expected string, got something else");
+    return NULL;
+  }
+
+  int64_t size = nbdkit_parse_size(s);
+  if (size == -1) {
+    PyErr_SetString (PyExc_ValueError, "Unable to parse string as size");
+    return NULL;
+  }
+
+  return PyLong_FromSize_t((size_t)size);
+}
+
 static PyMethodDef NbdkitMethods[] = {
   { "debug", debug, METH_VARARGS,
     "Print a debug message" },
   { "export_name", export_name, METH_VARARGS,
     "Return the optional export name negotiated with the client" },
+  { "parse_size", parse_size, METH_VARARGS,
+    "Parse human-readable size strings into bytes" },
   { "set_error", set_error, METH_VARARGS,
     "Store an errno value prior to throwing an exception" },
   { "shutdown", do_shutdown, METH_VARARGS,
diff --git a/plugins/python/nbdkit-python-plugin.pod b/plugins/python/nbdkit-python-plugin.pod
index 051b0237..ccc9406f 100644
--- a/plugins/python/nbdkit-python-plugin.pod
+++ b/plugins/python/nbdkit-python-plugin.pod
@@ -131,6 +131,11 @@ Record C<err> as the reason you are about to throw an exception. C<err>
 should correspond to usual errno values, where it may help to
 C<import errno>.
 
+=head3 C<nbdkit.parse_size(str)>
+
+Parse a string (such as "100M") into a size in bytes. Wraps the
+C<nbdkit_parse_size()> C function.
+
 =head3 C<nbdkit.shutdown()>
 
 Request asynchronous server shutdown.
diff --git a/tests/test-python-plugin.py b/tests/test-python-plugin.py
index 0b34d532..d4f379fc 100644
--- a/tests/test-python-plugin.py
+++ b/tests/test-python-plugin.py
@@ -34,12 +34,31 @@
 import nbdkit
 import pickle
 import base64
+import unittest
 
 API_VERSION = 2
 
 cfg = {}
 
 
+# Not nice, but there doesn't seem to be a better way of putting this
+class TestAPI(unittest.TestCase):
+
+    def test_parse_size(self):
+        self.assertEqual(nbdkit.parse_size('511'), 511)
+        self.assertEqual(nbdkit.parse_size('7k'), 7*1024)
+        self.assertEqual(nbdkit.parse_size('17M'), 17*1024*1024)
+
+        with self.assertRaises(TypeError):
+            nbdkit.parse_size(17)
+
+        with self.assertRaises(ValueError):
+            nbdkit.parse_size('foo')
+
+
+TestAPI().test_parse_size()
+
+
 def config(k, v):
     global cfg
     if k == "cfg":
-- 
2.31.1