[df] Avoid side-effects if external package interacts with distribute…

bindings/distrdf/python/DistRDF/Backends/__init__.py

@@ -16,12 +16,13 @@
 import types
 
 
-def build_backends_submodules(parentmodule: types.ModuleType) -> types.ModuleType:
+def build_backends_submodules(parentmodule: types.ModuleType, experimental: bool) -> types.ModuleType:
     """
     Helper function to create the submodules of the backends.
     """
-    for _, module_name, is_pkg in pkgutil.walk_packages(__path__):
+    from ROOT._distrdf import _raise_warning_if_experimental
 
+    for _, module_name, is_pkg in pkgutil.walk_packages(__path__):
         if is_pkg:
             # The actual python package with the backend implementation
             actual = importlib.import_module(__name__ + "." + module_name)
@@ -37,7 +38,7 @@ def build_backends_submodules(parentmodule: types.ModuleType) -> types.ModuleTyp
             dummy.__package__ = parentmodule
 
             # Attached functions
-            dummy.RDataFrame = actual.RDataFrame
+            dummy.RDataFrame = _raise_warning_if_experimental(actual.RDataFrame, experimental)
 
             setattr(parentmodule, module_name, dummy)
 

bindings/distrdf/python/DistRDF/__init__.py

@@ -12,12 +12,13 @@
 from __future__ import annotations
 
 import concurrent.futures
+import functools
 import importlib.util
 import logging
 import textwrap
 import types
 import warnings
-from typing import TYPE_CHECKING, Iterable
+from typing import TYPE_CHECKING, Callable, Iterable
 
 from .Backends import build_backends_submodules
 from .LiveVisualize import LiveVisualize
@@ -208,29 +209,33 @@ def FromSpec(jsonfile: str, *args, **kwargs) -> RDataFrame:
     )
 
 
-class _DeprecatedModule(types.ModuleType):
-    """A simple module type to raise a warning before usage."""
-
-    def __getattribute__(self, name):
-        msg_warng = textwrap.dedent(
-            """
-            In ROOT 6.36, the ROOT.RDF.Experimental.Distributed module has become just ROOT.RDF.Distributed. In the
-            future, the 'Experimental' keyword will be removed, so it is suggested to move to the stable API in user 
-            code. You can now change lines such as:
-            ```
-            connection = ... # your distributed Dask client or SparkContext
-            RDataFrame = ROOT.RDF.Experimental.Distributed.[Backend].RDataFrame
-            df = RDataFrame(..., [daskclient,sparkcontext] = connection)
-            ```
-            to simply:
-            ```
-            connection = ... # your distributed Dask client or SparkContext
-            df = ROOT.RDataFrame(..., executor = connection)
-            ```
-            """
-        )
-        warnings.warn(msg_warng, FutureWarning)
-        return super().__getattribute__(name)
+def _raise_warning_if_experimental(func: Callable, experimental: bool):
+    """A simple function decorator to raise a warning before usage."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        if experimental:
+            msg_warng = textwrap.dedent(
+                """
+                In ROOT 6.36, the ROOT.RDF.Experimental.Distributed module has become just ROOT.RDF.Distributed. In the
+                future, the 'Experimental' keyword will be removed, so it is suggested to move to the stable API in user 
+                code. You can now change lines such as:
+                ```
+                connection = ... # your distributed Dask client or SparkContext
+                RDataFrame = ROOT.RDF.Experimental.Distributed.[Backend].RDataFrame
+                df = RDataFrame(..., [daskclient,sparkcontext] = connection)
+                ```
+                to simply:
+                ```
+                connection = ... # your distributed Dask client or SparkContext
+                df = ROOT.RDataFrame(..., executor = connection)
+                ```
+                """
+            )
+            warnings.warn(msg_warng, FutureWarning, stacklevel=2)
+        return func(*args, **kwargs)
+
+    return wrapper
 
 
 def create_distributed_module(parentmodule, experimental: bool = False):
@@ -249,21 +254,18 @@ def create_distributed_module(parentmodule, experimental: bool = False):
     # distributed.__loader__ is not defined
     distributed.__package__ = parentmodule
 
-    distributed = build_backends_submodules(distributed)
+    distributed = build_backends_submodules(distributed, experimental)
 
     # Inject top-level functions
-    distributed.initialize = initialize
-    distributed.RunGraphs = RunGraphs
-    distributed.VariationsFor = VariationsFor
-    distributed.LiveVisualize = LiveVisualize
-    distributed.DistributeHeaders = DistributeHeaders
-    distributed.DistributeFiles = DistributeFiles
-    distributed.DistributeSharedLibs = DistributeSharedLibs
-    distributed.DistributeCppCode = DistributeCppCode
-    distributed.FromSpec = FromSpec
-
-    if experimental:
-        distributed.__class__ = _DeprecatedModule
+    distributed.initialize = _raise_warning_if_experimental(initialize, experimental)
+    distributed.RunGraphs = _raise_warning_if_experimental(RunGraphs, experimental)
+    distributed.VariationsFor = _raise_warning_if_experimental(VariationsFor, experimental)
+    distributed.LiveVisualize = _raise_warning_if_experimental(LiveVisualize, experimental)
+    distributed.DistributeHeaders = _raise_warning_if_experimental(DistributeHeaders, experimental)
+    distributed.DistributeFiles = _raise_warning_if_experimental(DistributeFiles, experimental)
+    distributed.DistributeSharedLibs = _raise_warning_if_experimental(DistributeSharedLibs, experimental)
+    distributed.DistributeCppCode = _raise_warning_if_experimental(DistributeCppCode, experimental)
+    distributed.FromSpec = _raise_warning_if_experimental(FromSpec, experimental)
 
     return distributed
 

roottest/python/distrdf/backends/check_explicit_api.py

@@ -90,6 +90,24 @@ def test_redefine_one_column(self, payload):
 class TestDeprecation:
     """Test deprecation warnings in the distributed module."""
 
+    MODULE_WARN = textwrap.dedent(
+        """
+    In ROOT 6.36, the ROOT.RDF.Experimental.Distributed module has become just ROOT.RDF.Distributed. In the
+    future, the 'Experimental' keyword will be removed, so it is suggested to move to the stable API in user 
+    code. You can now change lines such as:
+    ```
+    connection = ... # your distributed Dask client or SparkContext
+    RDataFrame = ROOT.RDF.Experimental.Distributed.[Backend].RDataFrame
+    df = RDataFrame(..., [daskclient,sparkcontext] = connection)
+    ```
+    to simply:
+    ```
+    connection = ... # your distributed Dask client or SparkContext
+    df = ROOT.RDataFrame(..., executor = connection)
+    ```
+    """
+    )
+
     def test_experimental_module_deprecation(self, payload):
         connection, backend = payload
         with pytest.warns() as warninglist:
@@ -100,50 +118,34 @@ def test_experimental_module_deprecation(self, payload):
                 RDataFrame = ROOT.RDF.Experimental.Distributed.Spark.RDataFrame
                 df = RDataFrame(10, npartitions=2, executor=connection)
 
-            msg_warng = textwrap.dedent(
-                """
-                In ROOT 6.36, the ROOT.RDF.Experimental.Distributed module has become just ROOT.RDF.Distributed. In the
-                future, the 'Experimental' keyword will be removed, so it is suggested to move to the stable API in user 
-                code. You can now change lines such as:
-                ```
-                connection = ... # your distributed Dask client or SparkContext
-                RDataFrame = ROOT.RDF.Experimental.Distributed.[Backend].RDataFrame
-                df = RDataFrame(..., [daskclient,sparkcontext] = connection)
-                ```
-                to simply:
-                ```
-                connection = ... # your distributed Dask client or SparkContext
-                df = ROOT.RDataFrame(..., executor = connection)
-                ```
-                """
-            )
             assert len(warninglist) == 1, f"{warninglist}"
             assert issubclass(warninglist[0].category, FutureWarning)
-            assert str(warninglist[0].message == msg_warng)
+            assert str(warninglist[0].message == self.MODULE_WARN)
 
             assert df.Count().GetValue() == 10
 
     def test_backend_argument_deprecation(self, payload):
         connection, backend = payload
-
-        if backend == "dask":
-            with pytest.warns():  # This will trigger the warning about experimental module deprecation tested above
+        with pytest.warns() as warninglist:
+            if backend == "dask":
                 RDataFrame = ROOT.RDF.Experimental.Distributed.Dask.RDataFrame
-            with pytest.warns(
-                FutureWarning,
-                match="The keyword argument 'daskclient' is not necessary anymore and will be removed in a future release",
-            ):
                 df = RDataFrame(10, npartitions=2, daskclient=connection)
-                assert df.Count().GetValue() == 10
-        elif backend == "spark":
-            with pytest.warns():  # This will trigger the warning about experimental module deprecation tested above
+            elif backend == "spark":
                 RDataFrame = ROOT.RDF.Experimental.Distributed.Spark.RDataFrame
-            with pytest.warns(
-                FutureWarning,
-                match="The keyword argument 'sparkcontext' is not necessary anymore and will be removed in a future release",
-            ):
                 df = RDataFrame(10, npartitions=2, sparkcontext=connection)
-                assert df.Count().GetValue() == 10
+
+            keyword = "daskclient" if backend == "dask" else "sparkcontext"
+            keyword_warn = (
+                f"The keyword argument '{keyword}' is not necessary anymore and will be removed in a future release"
+            )
+
+            assert len(warninglist) == 2, f"{warninglist}"
+            assert issubclass(warninglist[0].category, FutureWarning)
+            assert str(warninglist[0].message == self.MODULE_WARN)
+            assert issubclass(warninglist[1].category, FutureWarning)
+            assert str(warninglist[1].message == keyword_warn)
+
+            assert df.Count().GetValue() == 10
 
 
 if __name__ == "__main__":