Remove DataSlice._add_method, so that it is not accessible in public API.

PiperOrigin-RevId: 718389025
Change-Id: I74dd3e0e799155802a0275586ba60cc9389c8bd3

Author: mprekajski

py/koladata/expr/view_test.py

@@ -627,9 +627,7 @@ def test_data_slice_attrs_are_in_view(self):
         'internal_as_dense_array',
         'internal_as_py',
         'internal_is_any_schema',
-        'internal_is_compliant_attr_name',
         'internal_is_itemid_schema',
-        'internal_register_reserved_class_method_name',
         'is_mutable',
         'qtype',
         'set_attr',

py/koladata/types/data_item.py

@@ -19,8 +19,6 @@
 from arolla import arolla
 from koladata.expr import py_expr_eval_py_ext as _py_expr_eval_py_ext
 from koladata.types import data_item_py_ext as _data_item_py_ext
-# Used to initialize DataSlice, so it is available when defining subclasses of
-# DataItem.
 from koladata.types import data_slice
 
 _eval_op = _py_expr_eval_py_ext.eval_op
@@ -29,15 +27,15 @@
 
 
 ### Implementation of the DataItem's additional functionality.
-@DataItem._add_method('__hash__')  # pylint: disable=protected-access
+@data_slice.add_method(DataItem, '__hash__')  # pylint: disable=protected-access
 def _hash(self) -> int:
   return hash(self.fingerprint)
 
 
 # Ideally we'd do this only for functors, but we don't have a fast way
 # to check if a DataItem is a functor now. Note that SchemaItem overrides
 # this behavior.
-@DataItem._add_method('__call__')  # pylint: disable=protected-access
+@data_slice.add_method(DataItem, '__call__')  # pylint: disable=protected-access
 def _call(
     self, *args: Any, return_type_as: Any = data_slice.DataSlice, **kwargs: Any
 ) -> data_slice.DataSlice:
@@ -68,7 +66,7 @@ def register_bind_method_implementation(
 
 # Ideally we'd do this only for functors, but we don't have a fast way
 # to check if a DataItem is a functor now.
-@DataItem._add_method('bind')  # pylint: disable=protected-access
+@data_slice.add_method(DataItem, 'bind')  # pylint: disable=protected-access
 def bind(
     self,
     *,

py/koladata/types/data_item_test.py

@@ -124,21 +124,6 @@ def test_bag(self):
     x = x.with_bag(db)
     self.assertIsNotNone(x.get_bag())
 
-  def test_add_method(self):
-    # We are adding a method to DataSlice and should be visible in DataItem
-    # after the fact.
-    self.assertFalse(hasattr(data_item.DataItem, 'foo'))
-
-    @data_slice.DataSlice._add_method('foo')
-    def foo(self):
-      """Converts DataSlice to DenseArray."""
-      return self.internal_as_dense_array()
-
-    self.assertTrue(hasattr(data_item.DataItem, 'foo'))
-
-    x = ds('abc')
-    arolla.testing.assert_qvalue_allequal(x.foo(), arolla.dense_array(['abc']))
-
   def test_set_get_attr(self):
     db = data_bag.DataBag.empty()
     x = db.new(abc=ds(3.14))

py/koladata/types/data_slice.cc

@@ -312,19 +312,6 @@ bool IsCompliantAttrName(absl::string_view attr_name) {
              attr_name);
 }
 
-// classmethod
-absl::Nullable<PyObject*> PyDataSlice_is_compliant_attr_name(
-    PyObject* cls, PyObject* attr_name) {
-  arolla::python::DCheckPyGIL();
-  Py_ssize_t size;
-  const char* attr_name_ptr = PyUnicode_AsUTF8AndSize(attr_name, &size);
-  if (attr_name_ptr == nullptr) {
-    return nullptr;
-  }
-  auto attr_name_view = absl::string_view(attr_name_ptr, size);
-  return PyBool_FromLong(IsCompliantAttrName(attr_name_view));
-}
-
 absl::Nullable<PyObject*> PyDataSlice_getattro(PyObject* self,
                                                PyObject* attr_name) {
   arolla::python::DCheckPyGIL();
@@ -963,28 +950,6 @@ absl::Nullable<PyObject*> PyDataSlice_debug_repr(PyObject* self) {
       debug_repr.c_str(), static_cast<Py_ssize_t>(debug_repr.size()));
 }
 
-// classmethod
-absl::Nullable<PyObject*>
-PyDataSlice_internal_register_reserved_class_method_name(
-    PyTypeObject* cls, PyObject* method_name) {
-  arolla::python::DCheckPyGIL();
-  if (!PyUnicode_Check(method_name)) {
-    PyErr_SetString(PyExc_TypeError, "method name must be a string");
-    return nullptr;
-  }
-  Py_ssize_t size;
-  const char* method_name_ptr = PyUnicode_AsUTF8AndSize(method_name, &size);
-  if (method_name_ptr == nullptr) {
-    return nullptr;
-  }
-  auto method_name_view = absl::string_view(method_name_ptr, size);
-  if (size == 0 || method_name_view[0] != '_') {
-    PyDataSlice_GetReservedAttrsWithoutLeadingUnderscore().insert(
-        method_name_view);
-  }
-  Py_RETURN_NONE;
-}
-
 PyMethodDef kPyDataSlice_methods[] = {
     {"get_bag", PyDataSlice_get_bag, METH_NOARGS,
      "get_bag()\n"
@@ -1220,29 +1185,12 @@ schema. In case of ANY (or primitives), an empty list is returned.
      "Returns a format representation with a special support for non empty "
      "specification.\n\nDataSlice will be replaced with base64 encoded "
      "DataSlice.\nMust be used with kd.fstr or kde.fstr."},
-    {"internal_register_reserved_class_method_name",
-     (PyCFunction)PyDataSlice_internal_register_reserved_class_method_name,
-     METH_CLASS | METH_O,
-     "internal_register_reserved_class_method_name(method_name, /)\n"
-     "--\n\n"
-     "Registers a name to be reserved as a method of the DataSlice class.\n"
-     "\n"
-     "You must call this when adding new methods to the class in Python.\n"
-     "\n"
-     "Args:\n"
-     "  method_name: (str)\n"},
     {"_repr_with_params", (PyCFunction)PyDataSlice_repr_with_params,
      METH_FASTCALL | METH_KEYWORDS,
      "_repr_with_params("
      "*, depth=1, unbounded_type_max_len=-1, format_html=False)\n"
      "--\n\n"
      "Used to generate str representation for interactive repr in Colab."},
-    {"internal_is_compliant_attr_name",
-     (PyCFunction)PyDataSlice_is_compliant_attr_name, METH_CLASS | METH_O,
-     "internal_is_compliant_attr_name(attr_name, /)\n"
-     "--\n\n"
-     "Returns true iff `attr_name` can be accessed through "
-     "`getattr(slice, attr_name)`."},
     {"_debug_repr", (PyCFunction)PyDataSlice_debug_repr, METH_NOARGS,
      "_debug_repr()\n"
      "--\n\n"
@@ -1349,4 +1297,37 @@ PyTypeObject* PyDataSlice_Type() {
   return type;
 }
 
+// data_slice_py_ext module methods
+absl::Nullable<PyObject*> PyDataSliceModule_is_compliant_attr_name(
+    PyObject* /*module*/, PyObject* attr_name) {
+  arolla::python::DCheckPyGIL();
+  Py_ssize_t size;
+  const char* attr_name_ptr = PyUnicode_AsUTF8AndSize(attr_name, &size);
+  if (attr_name_ptr == nullptr) {
+    return nullptr;
+  }
+  auto attr_name_view = absl::string_view(attr_name_ptr, size);
+  return PyBool_FromLong(IsCompliantAttrName(attr_name_view));
+}
+
+absl::Nullable<PyObject*> PyDataSliceModule_register_reserved_class_method_name(
+    PyObject* /*module*/, PyObject* method_name) {
+  arolla::python::DCheckPyGIL();
+  if (!PyUnicode_Check(method_name)) {
+    PyErr_SetString(PyExc_TypeError, "method name must be a string");
+    return nullptr;
+  }
+  Py_ssize_t size;
+  const char* method_name_ptr = PyUnicode_AsUTF8AndSize(method_name, &size);
+  if (method_name_ptr == nullptr) {
+    return nullptr;
+  }
+  auto method_name_view = absl::string_view(method_name_ptr, size);
+  if (size == 0 || method_name_view[0] != '_') {
+    PyDataSlice_GetReservedAttrsWithoutLeadingUnderscore().insert(
+        method_name_view);
+  }
+  Py_RETURN_NONE;
+}
+
 }  // namespace koladata::python

py/koladata/types/data_slice.h

@@ -17,10 +17,25 @@
 
 #include <Python.h>
 
+#include "absl/base/nullability.h"
+
 namespace koladata::python {
 
 // Returns a PyType of Python DataSlice.
-PyTypeObject* PyDataSlice_Type();
+absl::Nullable<PyTypeObject*> PyDataSlice_Type();
+
+// Returns true iff `attr_name` can be accessed through:
+//   getattr(slice, attr_name)
+absl::Nullable<PyObject*> PyDataSliceModule_is_compliant_attr_name(
+    PyObject* /*module*/, PyObject* attr_name);
+
+// Registers a name in method_name to be reserved as a method of the DataSlice
+// class or it subclasses. For these registered method names, getattr will
+// invoke a PyObject_GenericGetAttr.
+//
+// This must be called when adding new methods to the class in Python.
+absl::Nullable<PyObject*> PyDataSliceModule_register_reserved_class_method_name(
+    PyObject* /*module*/, PyObject* method_name);
 
 }  // namespace koladata::python