Add midlog to polars and a clone option to midlog

Fix pandas midlog (endlog() returns a bare frame, so it must restart via
.raffa.startlog()), add the equivalent midlog to the polars Series and
DataFrame accessors, and give midlog an optional clone argument forwarded
to the restarted startlog. Add tests covering both backends.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: raffaelemancuso

README.md

@@ -147,7 +147,7 @@ See the [Examples](https://raffalib-python.readthedocs.io) page for the full wal
 | Module                  | What it provides                                                        |
 | ----------------------- | ---------------------------------------------------------------------- |
 | `raffalib.pandas`       | `.raffa` accessor: `startlog`/`endlog`/`midlog`, `join`, `freq`, `to_docx`, `add_prefix_if_not_exists`, `get_duplicates`, `sort_columns` |
-| `raffalib.polars`       | `.raffa` namespace: logging, `freq`, `crosstab`, `join`, `replace_string_with_null`, `to_docx` |
+| `raffalib.polars`       | `.raffa` namespace: `startlog`/`endlog`/`midlog`, `freq`, `crosstab`, `join`, `replace_string_with_null`, `to_docx` |
 | `raffalib.logging`      | `create_logger` — opinionated logging setup (plain or `rich`)          |
 | `raffalib.export_docx`  | `DocxFile` — low-level Word document/table builder                     |
 | `raffalib.tqdm`         | `tqdm_batch` — batched progress bars                                    |

docs/source/examples.rst

@@ -32,9 +32,11 @@ how many cell values changed.
    appends the elapsed time since ``startlog``. The signature is identical for
    pandas and polars.
 
-``midlog(custom_msg=None, timeit=True)`` *(pandas only)*
+``midlog(custom_msg=None, timeit=True, clone=False)``
    A convenience alias for ``endlog().startlog()`` — log the current step and
-   immediately start the next one.
+   immediately start the next one. Available on both pandas and polars. ``clone``
+   is forwarded to the restarted ``startlog``, so pass ``clone=True`` when the
+   next segment should detect value-level changes.
 
 Messages are emitted through the standard :mod:`logging` module. Configure it
 once with :func:`raffalib.create_logger` so the messages reach your console:

src/raffalib/pandas.py

@@ -73,11 +73,22 @@ def startlog(self, clone=False):
             self._series._initial_series = None
         return self._series
 
-    def midlog(self, custom_msg: str | None = None, timeit: bool = True):
+    def midlog(
+        self, custom_msg: str | None = None, timeit: bool = True, clone: bool = False
+    ):
         """
-        Alias for `.endlog().startlog()`
+        Alias for `.endlog().startlog(clone=clone)`.
+
+        :param custom_msg: A custom message to log before the actual log
+        :type custom_msg: str | None
+        :param timeit: Log the time it took for the operation
+        :type timeit: bool
+        :param clone: Forwarded to the restarted ``startlog``: clone the data so
+            the next segment can detect value-level changes when the shape is
+            unchanged.
+        :type clone: bool
         """
-        return self.endlog(custom_msg, timeit=timeit).startlog()
+        return self.endlog(custom_msg, timeit=timeit).raffa.startlog(clone=clone)
 
     def endlog(self, custom_msg: str | None = None, timeit: bool = True):
         """
@@ -179,12 +190,24 @@ def startlog(self, clone: bool = False) -> pd.DataFrame:
         return self._df
 
     def midlog(
-        self, custom_msg: str | None = None, timeit: bool = True
+        self,
+        custom_msg: str | None = None,
+        timeit: bool = True,
+        clone: bool = False,
     ) -> pd.DataFrame:
         """
-        Alias for `.endlog().startlog()`
+        Alias for `.endlog().startlog(clone=clone)`.
+
+        :param custom_msg: A custom message to log before the actual log
+        :type custom_msg: str | None
+        :param timeit: Log the time it took for the operation
+        :type timeit: bool
+        :param clone: Forwarded to the restarted ``startlog``: clone the data so
+            the next segment can detect value-level changes when the shape is
+            unchanged.
+        :type clone: bool
         """
-        return self.endlog(custom_msg, timeit=timeit).startlog()
+        return self.endlog(custom_msg, timeit=timeit).raffa.startlog(clone=clone)
 
     def endlog(
         self, custom_msg: str | None = None, timeit: bool = True

src/raffalib/polars.py

@@ -62,8 +62,8 @@ class RaffaPolarsSeriesUtils:
     The ``.raffa`` namespace on a :class:`polars.Series`.
 
     Registered automatically when :mod:`raffalib.polars` is imported. Provides
-    STATA-like change logging (:meth:`startlog` / :meth:`endlog`) and tabulation
-    helpers (:meth:`freq`, :meth:`crosstab`).
+    STATA-like change logging (:meth:`startlog` / :meth:`endlog` / :meth:`midlog`)
+    and tabulation helpers (:meth:`freq`, :meth:`crosstab`).
     """
 
     def __init__(self, series: pl.Series):
@@ -125,6 +125,23 @@ def endlog(self, custom_msg: str | None = None, timeit: bool = True) -> pl.Serie
         logger.info(msg)
         return self._series
 
+    def midlog(
+        self, custom_msg: str | None = None, timeit: bool = True, clone: bool = False
+    ) -> pl.Series:
+        """
+        Alias for `.endlog().startlog(clone=clone)`.
+
+        :param custom_msg: A custom message to log before the actual log
+        :type custom_msg: str | None
+        :param timeit: Log the time it took for the operation
+        :type timeit: bool
+        :param clone: Forwarded to the restarted ``startlog``: clone the data so
+            the next segment can detect value-level changes when the shape is
+            unchanged.
+        :type clone: bool
+        """
+        return self.endlog(custom_msg, timeit=timeit).raffa.startlog(clone=clone)
+
     def freq(self) -> pl.DataFrame:
         """
         Frequency table for the Series.
@@ -183,9 +200,9 @@ class RaffaPolarsDataFrameUtils:
     The ``.raffa`` namespace on a :class:`polars.DataFrame`.
 
     Registered automatically when :mod:`raffalib.polars` is imported. Provides
-    STATA-like change logging (:meth:`startlog` / :meth:`endlog`), a logging
-    :meth:`join` wrapper, tabulation helpers (:meth:`freq`, :meth:`crosstab`),
-    and :meth:`to_docx` export.
+    STATA-like change logging (:meth:`startlog` / :meth:`endlog` / :meth:`midlog`),
+    a logging :meth:`join` wrapper, tabulation helpers (:meth:`freq`,
+    :meth:`crosstab`), and :meth:`to_docx` export.
     """
 
     def __init__(self, df: pl.DataFrame):
@@ -254,6 +271,26 @@ def endlog(
         logger.info(msg)
         return self._df
 
+    def midlog(
+        self,
+        custom_msg: str | None = None,
+        timeit: bool = True,
+        clone: bool = False,
+    ) -> pl.DataFrame:
+        """
+        Alias for `.endlog().startlog(clone=clone)`.
+
+        :param custom_msg: A custom message to log before the actual log
+        :type custom_msg: str | None
+        :param timeit: Log the time it took for the operation
+        :type timeit: bool
+        :param clone: Forwarded to the restarted ``startlog``: clone the data so
+            the next segment can detect value-level changes when the shape is
+            unchanged.
+        :type clone: bool
+        """
+        return self.endlog(custom_msg, timeit=timeit).raffa.startlog(clone=clone)
+
     def replace_string_with_null(self, s):
         """
         Replace a sentinel string with null across all string columns.

tests/test_pandas.py

@@ -82,6 +82,31 @@ def test_series_endlog_removed_values(caplog):
     )
 
 
+def test_midlog_logs_then_restarts(caplog):
+    df = pd.DataFrame({"a": [1, 2, 3, 4]})
+    with caplog.at_level(logging.INFO, logger="raffalib.pandas"):
+        df.raffa.startlog()
+        df = df[df["a"] > 1]
+        df = df.raffa.midlog(timeit=False)  # logs the first step, restarts logging
+        df = df[df["a"] > 2]
+        df.raffa.endlog(timeit=False)
+    msgs = [r.message for r in caplog.records]
+    assert any("Removed 1/4 (25.00%) rows. New shape: (3, 1)." in m for m in msgs)
+    assert any("Removed 1/3 (33.33%) rows. New shape: (2, 1)." in m for m in msgs)
+
+
+def test_midlog_clone_enables_value_diff(caplog):
+    df = pd.DataFrame({"a": [1, 2, 3, 4]})
+    with caplog.at_level(logging.INFO, logger="raffalib.pandas"):
+        df.raffa.startlog()
+        df = df[df["a"] > 1]
+        df = df.raffa.midlog(timeit=False, clone=True)  # restart WITH a clone
+        df = df.assign(a=[99, 3, 4])
+        df.raffa.endlog(timeit=False)
+    msgs = [r.message for r in caplog.records]
+    assert any("Changed 1/3 (33.33%) values." in m for m in msgs)
+
+
 def test_join_inner_drops_source_columns(caplog):
     left = pd.DataFrame({"k": [1, 2, 3], "lv": [10, 20, 30]})
     right = pd.DataFrame({"k": [2, 3, 4], "rv": [200, 300, 400]})

tests/test_polars.py

@@ -58,6 +58,31 @@ def test_series_endlog_matches_pandas_format(caplog):
     )
 
 
+def test_midlog_logs_then_restarts(caplog):
+    df = pl.DataFrame({"a": [1, 2, 3, 4]})
+    with caplog.at_level(logging.INFO, logger="raffalib.polars"):
+        df.raffa.startlog()
+        df = df.filter(pl.col("a") > 1)
+        df = df.raffa.midlog(timeit=False)  # logs the first step, restarts logging
+        df = df.filter(pl.col("a") > 2)
+        df.raffa.endlog(timeit=False)
+    msgs = [r.message for r in caplog.records]
+    assert any("Removed 1/4 (25.00%) rows. New shape: (3, 1)." in m for m in msgs)
+    assert any("Removed 1/3 (33.33%) rows. New shape: (2, 1)." in m for m in msgs)
+
+
+def test_midlog_clone_enables_value_diff(caplog):
+    df = pl.DataFrame({"a": [1, 2, 3, 4]})
+    with caplog.at_level(logging.INFO, logger="raffalib.polars"):
+        df.raffa.startlog()
+        df = df.filter(pl.col("a") > 1)
+        df = df.raffa.midlog(timeit=False, clone=True)  # restart WITH a clone
+        df = df.with_columns(pl.Series("a", [99, 3, 4]))
+        df.raffa.endlog(timeit=False)
+    msgs = [r.message for r in caplog.records]
+    assert any("Changed 1/3 (33.33%) values." in m for m in msgs)
+
+
 def test_endlog_timeit_appends_duration(caplog):
     df = pl.DataFrame({"a": [1, 2, 3]})
     with caplog.at_level(logging.INFO, logger="raffalib.polars"):