From b010bb40824c2ac50feb7d50da2427fb45031c46 Mon Sep 17 00:00:00 2001
From: paul fisher <paul.fisher@tiledb.com>
Date: Mon, 13 Feb 2023 12:03:47 -0500
Subject: [PATCH] [python] Enable chaining on certain write methods. (#123)

Returns `self` from a few methods to enable chaining:

- DataFrame/DenseNDArray/SparseNDArray.write
- Collection.set

This enables useful chaining:

    some_coll.set('x', ...).set('y', ...)

    my_df = DataFrame.create(...).write(...)
    do_stuff_with(my_df)
---
 python-spec/src/somacore/collection.py            |  3 ++-
 python-spec/src/somacore/data.py                  | 10 ++++++----
 python-spec/src/somacore/ephemeral/collections.py |  3 ++-
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/python-spec/src/somacore/collection.py b/python-spec/src/somacore/collection.py
index 145f0b9e..2a5358a3 100644
--- a/python-spec/src/somacore/collection.py
+++ b/python-spec/src/somacore/collection.py
@@ -184,7 +184,7 @@ def __setitem__(self, key: str, value: _Elem) -> None:
     @abc.abstractmethod
     def set(
         self, key: str, value: _Elem, *, use_relative_uri: Optional[bool] = None
-    ) -> None:
+    ) -> Self:
         """Sets an entry of this collection. [lifecycle: experimental]
 
         Important note: Because parent objects may need to share
@@ -207,6 +207,7 @@ def set(
             not share a relative URI base, or use of relative URIs is not
             possible at all, the collection should throw an exception.
             If ``False``, will always use an absolute URI.
+        :return: ``self``, to enable method chaining.
         """
         raise NotImplementedError()
 
diff --git a/python-spec/src/somacore/data.py b/python-spec/src/somacore/data.py
index 2ecdbbad..d47f1ea0 100644
--- a/python-spec/src/somacore/data.py
+++ b/python-spec/src/somacore/data.py
@@ -138,7 +138,7 @@ def write(
         values: Union[pa.RecordBatch, pa.Table],
         *,
         platform_config: Optional[options.PlatformConfig] = None,
-    ) -> None:
+    ) -> Self:
         """Writes the data from an Arrow table to the persistent object.
         [lifecycle: experimental]
 
@@ -148,6 +148,7 @@ def write(
         :param values: An Arrow table containing all columns, including
             the index columns. The schema for the values must match
             the schema for the ``DataFrame``.
+        :return: ``self``, to enable method chaining.
         """
         raise NotImplementedError()
 
@@ -288,7 +289,7 @@ def write(
         values: pa.Tensor,
         *,
         platform_config: Optional[options.PlatformConfig] = None,
-    ) -> None:
+    ) -> Self:
         """Writes an Arrow tensor to a subarray of the persistent object.
         [lifecycle: experimental]
 
@@ -300,6 +301,7 @@ def write(
             See :meth:`read` for details about indexing.
         :param values: The values to be written to the subarray.  Must have
             the same shape as ``coords``, and matching type to the array.
+        :return: ``self``, to enable method chaining.
         """
         raise NotImplementedError()
 
@@ -382,11 +384,12 @@ def write(
         values: SparseArrowData,
         *,
         platform_config: Optional[options.PlatformConfig] = None,
-    ) -> None:
+    ) -> Self:
         """Writes a Tensor to a subarray of the persistent object.
         [lifecycle: experimental]
 
         :param values: The values to write to the array.
+        :return: ``self``, to enable method chaining.
 
         **Value types:**
 
@@ -397,7 +400,6 @@ def write(
 
         Arrow table: a COO table, with columns named ``soma_dim_0``, ...,
         ``soma_dim_N`` and ``soma_data``, to be written to the array.
-
         """
         raise NotImplementedError()
 
diff --git a/python-spec/src/somacore/ephemeral/collections.py b/python-spec/src/somacore/ephemeral/collections.py
index 568831e4..ec9bb406 100644
--- a/python-spec/src/somacore/ephemeral/collections.py
+++ b/python-spec/src/somacore/ephemeral/collections.py
@@ -79,9 +79,10 @@ def mode(self) -> options.OpenMode:
 
     def set(
         self, key: str, value: _Elem, *, use_relative_uri: Optional[bool] = None
-    ) -> None:
+    ) -> Self:
         del use_relative_uri  # Ignored.
         self._entries[key] = value
+        return self
 
     def __getitem__(self, key: str) -> _Elem:
         return self._entries[key]