Add output_class argument to to_py. It will be actually implemented as a follow-up.

PiperOrigin-RevId: 875733422
Change-Id: Id8aa2246b9c7ce79837c4dc8c333edeec539e510

Author: atorero

docs/api_reference.md

@@ -10901,7 +10901,7 @@ Alias for [kd.proto.to_proto_bytes](#kd.proto.to_proto_bytes) operator.
 
 Alias for [kd.proto.to_proto_json](#kd.proto.to_proto_json) operator.
 
-### `kd.to_py(ds: DataSlice, max_depth: int = 2, obj_as_dict: bool = False, include_missing_attrs: bool = True) -> Any` {#kd.to_py}
+### `kd.to_py(ds: DataSlice, max_depth: int = 2, obj_as_dict: bool = False, include_missing_attrs: bool = True, output_class: type[Any] | None = None) -> Any` {#kd.to_py}
 
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Returns a readable python object from a DataSlice.
 
@@ -10915,7 +10915,8 @@ Args:
   obj_as_dict: Whether to convert objects to python dicts. By default objects
     are converted to automatically constructed &#39;Obj&#39; dataclass instances.
   include_missing_attrs: whether to include attributes with None value in
-    objects.</code></pre>
+    objects.
+  output_class: If not None, will be used recursively as the output type.</code></pre>
 
 ### `kd.to_pylist(x: DataSlice) -> list[Any]` {#kd.to_pylist}
 
@@ -13454,7 +13455,7 @@ Args:
 Returns:
   A new DataSlice with items selected by indices.</code></pre>
 
-### `DataSlice.to_py(ds: DataSlice, max_depth: int = 2, obj_as_dict: bool = False, include_missing_attrs: bool = True) -> Any` {#DataSlice.to_py}
+### `DataSlice.to_py(ds: DataSlice, max_depth: int = 2, obj_as_dict: bool = False, include_missing_attrs: bool = True, output_class: Any | None = None) -> Any` {#DataSlice.to_py}
 
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Returns a readable python object from a DataSlice.
 
@@ -13467,7 +13468,8 @@ Args:
   obj_as_dict: Whether to convert objects to python dicts. By default objects
     are converted to automatically constructed &#39;Obj&#39; dataclass instances.
   include_missing_attrs: whether to include attributes with None value in
-    objects.</code></pre>
+    objects.
+  output_class: If not None, will be used recursively as the output type.</code></pre>
 
 ### `DataSlice.to_pytree(ds: DataSlice, max_depth: int = 2, include_missing_attrs: bool = True) -> Any` {#DataSlice.to_pytree}
 

py/koladata/base/py_conversions/to_py.cc

@@ -98,12 +98,14 @@ class ToPyVisitor : internal::AbstractVisitor {
  public:
   ToPyVisitor(bool obj_as_dict, bool include_missing_attrs,
               const absl::flat_hash_set<ObjectId>& objects_not_to_convert,
-              DataBagPtr db, internal::DataBagImpl::FallbackSpan fallback_span)
+              DataBagPtr db, internal::DataBagImpl::FallbackSpan fallback_span,
+              PyObject* output_class)
       : obj_as_dict_(obj_as_dict),
         include_missing_attrs_(include_missing_attrs),
         objects_not_to_convert_(objects_not_to_convert),
         db_(std::move(db)),
-        fallback_span_(fallback_span) {}
+        fallback_span_(fallback_span),
+        output_class_(PyObjectPtr::NewRef(output_class)) {}
 
   ItemToPyConverter GetItemToPyConverter(const DataItem& schema) {
     return [&](const DataItem& item) -> absl::StatusOr<PyObjectPtr> {
@@ -408,26 +410,34 @@ class ToPyVisitor : internal::AbstractVisitor {
   DataClassesUtil dataclasses_util_;
   DataBagPtr db_;
   internal::DataBagImpl::FallbackSpan fallback_span_;
+  PyObjectPtr output_class_;
 };
 
 PyObject* absl_nullable ToPyImplInternal(
     const DataSlice& ds, DataBagPtr bag, bool obj_as_dict,
-    bool include_missing_attrs, const absl::flat_hash_set<ObjectId>& leaf_ids) {
+    bool include_missing_attrs, const absl::flat_hash_set<ObjectId>& leaf_ids,
+    PyObject* output_class) {
   if (ds.IsEmpty() || bag == nullptr ||
       GetNarrowedSchema(ds).is_primitive_schema()) {
     ASSIGN_OR_RETURN(PyObjectPtr res, PyObjectFromDataSlice(ds),
                      arolla::python::SetPyErrFromStatus(_));
     return res.release();
   }
+  if (obj_as_dict && output_class != Py_None) {
+    arolla::python::SetPyErrFromStatus(absl::InvalidArgumentError(
+        "obj_as_dict cannot be used with output_class"));
+    return nullptr;
+  }
   const DataSlice& schema = ds.GetSchema();
   DCHECK(bag != nullptr);
 
   FlattenFallbackFinder fb_finder(*bag);
   const internal::DataBagImpl::FallbackSpan fallback_span =
       fb_finder.GetFlattenFallbacks();
   // We use original DataBag in ToPyVisitor.
-  std::shared_ptr<ToPyVisitor> visitor = std::make_shared<ToPyVisitor>(
-      obj_as_dict, include_missing_attrs, leaf_ids, bag, fallback_span);
+  std::shared_ptr<ToPyVisitor> visitor =
+      std::make_shared<ToPyVisitor>(obj_as_dict, include_missing_attrs,
+                                    leaf_ids, bag, fallback_span, output_class);
   // We use extracted DataBag for traversal.
   internal::Traverser<ToPyVisitor> traverse_op(bag->GetImpl(), fallback_span,
                                                visitor);
@@ -451,7 +461,8 @@ PyObject* absl_nullable ToPyImplInternal(
 
 PyObject* absl_nullable ToPyImpl(const DataSlice& ds, DataBagPtr bag,
                                  int max_depth, bool obj_as_dict,
-                                 bool include_missing_attrs) {
+                                 bool include_missing_attrs,
+                                 PyObject* output_class) {
   // When `max_depth != -1`, we want objects/dicts/lists at `max_depth`
   // to be kept as DataItems and not converted to Python objects.
   // To do that, we extract a DataSlice, and keep track of the leaf DataItems.
@@ -474,10 +485,11 @@ PyObject* absl_nullable ToPyImpl(const DataSlice& ds, DataBagPtr bag,
             /*casting_callback=*/std::nullopt, std::move(leaf_callback)),
         arolla::python::SetPyErrFromStatus(_));
     return ToPyImplInternal(extracted_ds, ds.GetBag(), obj_as_dict,
-                            include_missing_attrs, objects_not_to_convert);
+                            include_missing_attrs, objects_not_to_convert,
+                            output_class);
   }
   return ToPyImplInternal(ds, nullptr, obj_as_dict, include_missing_attrs,
-                          objects_not_to_convert);
+                          objects_not_to_convert, output_class);
 }
 
 }  // namespace
@@ -487,9 +499,9 @@ PyObject* absl_nullable PyDataSlice_to_py(PyObject* self,
                                           Py_ssize_t nargs) {
   arolla::python::DCheckPyGIL();
   arolla::python::PyCancellationScope cancellation_scope;
-  if (nargs != 3) {
+  if (nargs != 4) {
     PyErr_Format(PyExc_ValueError,
-                 "DataSlice._to_py_impl accepts exactly 3 arguments, got %d",
+                 "DataSlice._to_py_impl accepts exactly 4 arguments, got %d",
                  nargs);
     return nullptr;
   }
@@ -518,9 +530,15 @@ PyObject* absl_nullable PyDataSlice_to_py(PyObject* self,
     return nullptr;
   }
   const bool include_missing_attrs = py_args[2] == Py_True;
+  PyObject* output_class = py_args[3];
+  if (output_class == nullptr) {
+    PyErr_Format(PyExc_TypeError,
+                 "expecting output_class to be a class, got nullptr");
+    return nullptr;
+  }
 
   return ToPyImpl(ds, ds.GetBag(), max_depth, obj_as_dict,
-                  include_missing_attrs);
+                  include_missing_attrs, output_class);
 }
 
 }  // namespace koladata::python

py/koladata/functions/py_conversions.py

@@ -70,6 +70,7 @@ def to_py(
     max_depth: int = 2,
     obj_as_dict: bool = False,
     include_missing_attrs: bool = True,
+    output_class: type[Any] | None = None,
 ) -> Any:
   """Returns a readable python object from a DataSlice.
 
@@ -84,9 +85,10 @@ def to_py(
       are converted to automatically constructed 'Obj' dataclass instances.
     include_missing_attrs: whether to include attributes with None value in
       objects.
+    output_class: If not None, will be used recursively as the output type.
   """
   return ds._to_py_impl(  # pylint: disable=protected-access
-      max_depth, obj_as_dict, include_missing_attrs
+      max_depth, obj_as_dict, include_missing_attrs, output_class
   )
 
 

py/koladata/types/data_slice.cc

@@ -1020,7 +1020,7 @@ If `schema` is set, that schema is used, otherwise the schema is inferred from
      "Returns an UNSPECIFIED with DataSlice QType."},
     {"_to_py_impl", (PyCFunction)PyDataSlice_to_py, METH_FASTCALL,
      "_to_py_impl(ds, /, max_depth=-1, obj_as_dict=False, "
-     "include_missing_attrs=True)\n"
+     "include_missing_attrs=True, output_class=None)\n"
      "--\n\n"
      "Returns a Python object equivalent to this DataSlice.\n"},
     {"internal_as_py", PyDataSlice_internal_as_py, METH_NOARGS,

py/koladata/types/data_slice.py

@@ -498,12 +498,14 @@ def _explode(
   return _eval_op('kd.explode', self, ndim)
 
 
+# TODO(b/483020237) make arguments keyword-only.
 @add_method(DataSlice, 'to_py')
 def to_py(
     ds: DataSlice,
     max_depth: int = 2,
     obj_as_dict: bool = False,
     include_missing_attrs: bool = True,
+    output_class: Any | None = None,
 ) -> Any:
   """Returns a readable python object from a DataSlice.
 
@@ -517,15 +519,18 @@ def to_py(
       are converted to automatically constructed 'Obj' dataclass instances.
     include_missing_attrs: whether to include attributes with None value in
       objects.
+    output_class: If not None, will be used recursively as the output type.
   """
   return ds._to_py_impl(  # pylint: disable=protected-access
-      max_depth, obj_as_dict, include_missing_attrs
+      max_depth, obj_as_dict, include_missing_attrs, output_class
   )
 
 
 @add_method(DataSlice, 'to_pytree')
 def to_pytree(
-    ds: DataSlice, max_depth: int = 2, include_missing_attrs: bool = True
+    ds: DataSlice,
+    max_depth: int = 2,
+    include_missing_attrs: bool = True,
 ) -> Any:
   """Returns a readable python object from a DataSlice.
 
@@ -542,7 +547,7 @@ def to_pytree(
       objects.
   """
   return ds._to_py_impl(  # pylint: disable=protected-access
-      max_depth, True, include_missing_attrs
+      max_depth, True, include_missing_attrs, None
   )
 
 

py/koladata/types/type_defs.pyi

@@ -152,13 +152,13 @@ class DataSlice(arolla.QValue):
     def with_name(self, name: Any) -> Self: ...
 
     def to_py(
-        self, max_depth: int = 2, obj_as_dict: bool = False, include_missing_attrs: bool = True
+        self, max_depth: int = 2, obj_as_dict: bool = False, include_missing_attrs: bool = True, output_class: Any | None = None
     ) -> Any: ...
     def to_pytree(
         self, max_depth: int = 2, include_missing_attrs: bool = True
     ) -> Any: ...
     def _to_py_impl(
-        self, max_depth: int = 2, obj_as_dict: bool = False, include_missing_attrs: bool = True
+        self, max_depth: int = 2, obj_as_dict: bool = False, include_missing_attrs: bool = True, output_class: Any | None = None
     ) -> Any: ...
     def _to_proto(self, message_class: Any) -> Any: ...
     def internal_as_py(self) -> Any: ...