Allocate scratch space when loading databases

docs/usage.md

@@ -176,9 +176,16 @@ with open('hs.db', 'wb') as f:
     f.write(serialized)
 
 # Deserializing (loading from bytes):
-db = hyperscan.loadb(serialized)
+db = hyperscan.loadb(serialized, mode=hyperscan.HS_MODE_BLOCK)
+db.scan(b'foobar', match_event_handler=callback)
 ```
 
+Use the same scanning mode as the serialized database. By default, `loadb`
+allocates scratch space, so the restored database can scan immediately.
+Pass `alloc_scratch=False` to manage scratch allocation yourself; assign
+`db.scratch = hyperscan.Scratch(db)` or pass an appropriate scratch object
+when scanning.
+
 ## Chimera Mode
 
 ```python

src/hyperscan/__init__.pyi

@@ -82,12 +82,14 @@ def dumpb(database: "Database") -> bytes:
 
     """
 
-def loadb(buf: ByteString, mode: int) -> "Database":
+def loadb(buf: ByteString, mode: int, alloc_scratch: bool = True) -> "Database":
     """Deserializes a Hyperscan database.
 
     Args:
         buf (bytes): A serialized Hyperscan database.
         mode (int): The expected mode of the database.
+        alloc_scratch (bool, optional): Allocate scratch space for scanning.
+            Defaults to True. If False, supply scratch space before scanning.
 
     Returns:
         :class:`Database`: The deserialized database instance.

src/hyperscan/extension.c

@@ -1465,24 +1465,40 @@ static PyObject *loadb(PyObject *self, PyObject *args, PyObject *kwds)
 
   char *buf;
   PyObject *obuf = Py_None;
-  PyObject *odb;
-  odb = PyObject_CallFunctionObjArgs((PyObject *)&DatabaseType, NULL);
-  Database *db = (Database *)odb;
+  unsigned int mode;
+  int alloc_scratch = 1;
 
-  static char *kwlist[] = {"buf", "mode", NULL};
-  if (!PyArg_ParseTupleAndKeywords(args, kwds, "OI", kwlist, &obuf, &db->mode))
+  static char *kwlist[] = {"buf", "mode", "alloc_scratch", NULL};
+  if (!PyArg_ParseTupleAndKeywords(
+        args, kwds, "OI|p", kwlist, &obuf, &mode, &alloc_scratch))
     HS_LOCK_RETURN_NULL();
   if (!PyBytes_Check(obuf)) {
     PyErr_SetString(PyExc_TypeError, "buf must be a bytestring");
     HS_LOCK_RETURN_NULL();
   }
 
+  PyObject *odb = PyObject_CallFunctionObjArgs((PyObject *)&DatabaseType, NULL);
+  if (odb == NULL)
+    HS_LOCK_RETURN_NULL();
+  Database *db = (Database *)odb;
+  db->mode = mode;
+
   Py_ssize_t length = PyBytes_Size(obuf);
   buf = PyBytes_AsString(obuf);
   hs_error_t err = hs_deserialize_database(buf, length, &db->hs_db);
-  HANDLE_HYPERSCAN_ERR(err, NULL);
-  if (PyErr_Occurred())
-    HS_LOCK_RETURN_NULL();
+  if (err != HS_SUCCESS) {
+    Py_DECREF(odb);
+    HANDLE_HYPERSCAN_ERR(err, NULL);
+  }
+  if (alloc_scratch) {
+    PyObject *scratch =
+      PyObject_CallFunction((PyObject *)&ScratchType, "O", odb);
+    if (scratch == NULL) {
+      Py_DECREF(odb);
+      HS_LOCK_RETURN_NULL();
+    }
+    db->scratch = scratch;
+  }
   HS_LOCK_RETURN(odb);
 }
 
@@ -1499,11 +1515,14 @@ static PyMethodDef HyperscanMethods[] = {
   {"loadb",
    (PyCFunction)loadb,
    METH_VARARGS | METH_KEYWORDS,
-   "loadb(buf, mode)\n"
+   "loadb(buf, mode, alloc_scratch=True)\n"
    "    Deserializes a Hyperscan database.\n\n"
    "    Args:\n"
-   "        buf (:obj:`bytearray`): A serialized Hyperscan database.\n"
-   "        mode (int): The expected mode of the database.\n\n"
+   "        buf (bytes): A serialized Hyperscan database.\n"
+   "        mode (int): The expected mode of the database.\n"
+   "        alloc_scratch (bool, optional): Allocate scratch space for\n"
+   "            scanning. Defaults to True. If False, supply scratch\n"
+   "            space before scanning.\n\n"
    "    Returns:\n"
    "        :class:`Database`: The deserialized database instance.\n\n"},
   {NULL}};

tests/test_hyperscan.py

@@ -386,7 +386,8 @@ def test_database_deserialize_scan(db_fixture_name, request, mocker):
     original_db: hyperscan.Database = request.getfixturevalue(db_fixture_name)
     serialized_db = hyperscan.dumpb(original_db)
     db = hyperscan.loadb(serialized_db, original_db.mode)
-    db.scratch = hyperscan.Scratch(db)
+    assert isinstance(db.scratch, hyperscan.Scratch)
+    assert db.scratch.database is db
     if db_fixture_name == "database_stream":
         with db.stream(match_event_handler=callback) as stream:
             stream.scan(b"foobar")
@@ -395,6 +396,46 @@ def test_database_deserialize_scan(db_fixture_name, request, mocker):
         if db_fixture_name == "database_vector":
             buf = [buf]
         db.scan(buf, match_event_handler=callback)
+    callback.assert_called()
+
+
+@pytest.mark.parametrize(
+    "db_fixture_name",
+    ["database_stream", "database_block", "database_vector"],
+)
+@pytest.mark.parametrize("alloc_scratch", [False, True])
+def test_database_deserialize_scratch(db_fixture_name, alloc_scratch, request, mocker):
+    original_db = request.getfixturevalue(db_fixture_name)
+    db = hyperscan.loadb(
+        hyperscan.dumpb(original_db), original_db.mode, alloc_scratch=alloc_scratch
+    )
+    if alloc_scratch:
+        assert isinstance(db.scratch, hyperscan.Scratch)
+    else:
+        assert db.scratch is None
+        db.scratch = hyperscan.Scratch(db)
+
+    callback = mocker.Mock(return_value=None)
+    if db_fixture_name == "database_stream":
+        with db.stream(match_event_handler=callback) as stream:
+            stream.scan(b"foobar")
+    else:
+        data = [b"foo", b"bar"] if db_fixture_name == "database_vector" else b"foobar"
+        db.scan(data, match_event_handler=callback)
+    callback.assert_called()
+
+
+@pytest.mark.parametrize(
+    "data, error",
+    [
+        (None, TypeError),
+        (b"", hyperscan.InvalidError),
+        (b"not a database", hyperscan.InvalidError),
+    ],
+)
+def test_database_deserialize_invalid(data, error):
+    with pytest.raises(error):
+        hyperscan.loadb(data, hyperscan.HS_MODE_BLOCK)
 
 
 def test_database_exception_in_callback(database_block, mocker):