Simplify the C interface of DataBag.adopt_stub

PiperOrigin-RevId: 721717166
Change-Id: I7502cc4caa98755250a3bc7976c4d0e3c9f502d0

Author: mprekajski

py/koladata/types/data_bag.cc

@@ -1207,18 +1207,10 @@ absl::Nullable<PyObject*> PyDataBag_merge_inplace(PyObject* self,
   Py_RETURN_NONE;
 }
 
-absl::Nullable<PyObject*> PyDataBag_adopt(PyObject* self,
-                                          PyObject* const* py_args,
-                                          Py_ssize_t nargs) {
+absl::Nullable<PyObject*> PyDataBag_adopt(PyObject* self, PyObject* ds) {
   arolla::python::DCheckPyGIL();
-  DataBagPtr db = UnsafeDataBagPtr(self);
-  if (nargs != 1) {
-    PyErr_Format(PyExc_ValueError,
-                 "DataBag.adopt accepts exactly 1 argument, got %d", nargs);
-    return nullptr;
-  }
-
-  const DataSlice* slice = UnwrapDataSlice(py_args[0], "slice");
+  const DataBagPtr& db = UnsafeDataBagPtr(self);
+  const DataSlice* slice = UnwrapDataSlice(ds, "slice");
   if (!slice) {
     return nullptr;
   }
@@ -1231,18 +1223,10 @@ absl::Nullable<PyObject*> PyDataBag_adopt(PyObject* self,
   return WrapPyDataSlice(slice->WithBag(std::move(db)));
 }
 
-absl::Nullable<PyObject*> PyDataBag_adopt_stub(PyObject* self,
-                                               PyObject* const* py_args,
-                                               Py_ssize_t nargs) {
+absl::Nullable<PyObject*> PyDataBag_adopt_stub(PyObject* self, PyObject* ds) {
   arolla::python::DCheckPyGIL();
-  DataBagPtr db = UnsafeDataBagPtr(self);
-  if (nargs != 1) {
-    PyErr_Format(PyExc_ValueError,
-                 "DataBag.adopt_stub accepts exactly 1 argument, got %d",
-                 nargs);
-    return nullptr;
-  }
-  const DataSlice* slice = UnwrapDataSlice(py_args[0], "slice");
+  const DataBagPtr& db = UnsafeDataBagPtr(self);
+  const DataSlice* slice = UnwrapDataSlice(ds, "slice");
   if (!slice) {
     return nullptr;
   }
@@ -1663,7 +1647,7 @@ have different ids.
      "*bags)\n"
      "--\n\n"
      "DataBag._merge_inplace"},
-    {"adopt", (PyCFunction)PyDataBag_adopt, METH_FASTCALL,
+    {"adopt", (PyCFunction)PyDataBag_adopt, METH_O,
      "adopt(slice, /)\n"
      "--\n\n"
      R"""(Adopts all data reachable from the given slice into this DataBag.
@@ -1674,7 +1658,7 @@ have different ids.
 Returns:
   The DataSlice with this DataBag (including adopted data) attached.
 )"""},
-    {"adopt_stub", (PyCFunction)PyDataBag_adopt_stub, METH_FASTCALL,
+    {"adopt_stub", (PyCFunction)PyDataBag_adopt_stub, METH_O,
      "adopt_stub(slice, /)\n"
      "--\n\n"
      R"""(Copies the given DataSlice's schema stub into this DataBag.

py/koladata/types/data_bag_test.py

@@ -1741,14 +1741,12 @@ def test_merge_inplace_not_databags(self):
       db1.merge_inplace([x1])
 
   def test_adopt_args_errors(self):
-    with self.assertRaisesWithLiteralMatch(
-        ValueError, 'DataBag.adopt accepts exactly 1 argument, got 0'
-    ):
+    with self.assertRaises(TypeError):  # Python runtime error.
       bag().adopt()
-    with self.assertRaisesWithLiteralMatch(
-        ValueError, 'DataBag.adopt accepts exactly 1 argument, got 2'
-    ):
+    with self.assertRaises(TypeError):  # Python runtime error.
       bag().adopt(ds(1), ds(2))
+    with self.assertRaisesRegex(TypeError, 'expecting slice to be a DataSlice'):
+      bag().adopt(bag())
 
   def test_adopt_immutable(self):
     db1 = bag()
@@ -1777,14 +1775,12 @@ def test_adopt(self):
     testing.assert_equivalent(o3.y.no_bag(), ds(2))
 
   def test_adopt_stub_args_errors(self):
-    with self.assertRaisesWithLiteralMatch(
-        ValueError, 'DataBag.adopt_stub accepts exactly 1 argument, got 0'
-    ):
+    with self.assertRaises(TypeError):  # Python runtime error.
       bag().adopt_stub()
-    with self.assertRaisesWithLiteralMatch(
-        ValueError, 'DataBag.adopt_stub accepts exactly 1 argument, got 2'
-    ):
+    with self.assertRaises(TypeError):  # Python runtime error.
       bag().adopt_stub(ds(1), ds(2))
+    with self.assertRaisesRegex(TypeError, 'expecting slice to be a DataSlice'):
+      bag().adopt_stub(bag())
 
   def test_adopt_stub_immutable(self):
     db1 = bag()