Add str method to datamodule

src/lightning/pytorch/core/datamodule.py

@@ -14,6 +14,7 @@
 """LightningDataModule for loading DataLoaders with ease."""
 
 import inspect
+import os
 from collections.abc import Iterable
 from typing import IO, Any, Optional, Union, cast
 
@@ -244,3 +245,34 @@ def load_from_checkpoint(
             **kwargs,
         )
         return cast(Self, loaded)
+
+    def __str__(self) -> str:
+        """Return a string representation of the datasets that are setup.
+
+        Returns:
+            A string representation of the datasets that are setup.
+
+        """
+        datasets_info: Optional[list[str]] = []
+
+        def len_implemented(obj: Dataset) -> bool:
+            try:
+                len(obj)
+                return True
+            except NotImplementedError:
+                return False
+
+        for attr_name in dir(self):
+            attr = getattr(self, attr_name)
+
+            # Get Dataset information
+            if isinstance(attr, Dataset):
+                if hasattr(attr, "__len__") and len_implemented(attr):
+                    datasets_info.append(f"name={attr_name}, size={len(attr)}")
+                else:
+                    datasets_info.append(f"name={attr_name}, size=Unavailable")
+
+        if not datasets_info:
+            return "No datasets are set up."
+
+        return os.linesep.join(datasets_info)

src/lightning/pytorch/demos/boring_classes.py

@@ -188,6 +188,38 @@ def predict_dataloader(self) -> DataLoader:
         return DataLoader(self.random_predict)
 
 
+class BoringDataModuleNoLen(LightningDataModule):
+    """
+    .. warning::  This is meant for testing/debugging and is experimental.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.random_full = RandomIterableDataset(32, 64 * 4)
+
+
+class BoringDataModuleLenNotImplemented(LightningDataModule):
+    """
+    .. warning::  This is meant for testing/debugging and is experimental.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        class DS(Dataset):
+            def __init__(self, size: int, length: int):
+                self.len = length
+                self.data = torch.randn(length, size)
+
+            def __getitem__(self, index: int) -> Tensor:
+                return self.data[index]
+
+            def __len__(self) -> int:
+                raise NotImplementedError
+
+        self.random_full = DS(32, 64 * 4)
+
+
 class ManualOptimBoringModel(BoringModel):
     """
     .. warning::  This is meant for testing/debugging and is experimental.

tests/tests_pytorch/core/test_datamodules.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
 import pickle
 from argparse import Namespace
 from dataclasses import dataclass
@@ -22,7 +23,12 @@
 import torch
 from lightning.pytorch import LightningDataModule, Trainer, seed_everything
 from lightning.pytorch.callbacks import ModelCheckpoint
-from lightning.pytorch.demos.boring_classes import BoringDataModule, BoringModel
+from lightning.pytorch.demos.boring_classes import (
+    BoringDataModule,
+    BoringDataModuleLenNotImplemented,
+    BoringDataModuleNoLen,
+    BoringModel,
+)
 from lightning.pytorch.profilers.simple import SimpleProfiler
 from lightning.pytorch.trainer.states import TrainerFn
 from lightning.pytorch.utilities import AttributeDict
@@ -510,3 +516,64 @@ def prepare_data(self):
         durations = profiler.recorded_durations[key]
         assert len(durations) == 1
         assert durations[0] > 0
+
+
+def test_datamodule_string_no_datasets():
+    dm = BoringDataModule()
+    del dm.random_full
+    expected_output = "No datasets are set up."
+    assert str(dm) == expected_output
+
+
+def test_datamodule_string_no_length():
+    dm = BoringDataModuleNoLen()
+    expected_output = "name=random_full, size=Unavailable"
+    assert str(dm) == expected_output
+
+
+def test_datamodule_string_length_not_implemented():
+    dm = BoringDataModuleLenNotImplemented()
+    expected_output = "name=random_full, size=Unavailable"
+    assert str(dm) == expected_output
+
+
+def test_datamodule_string_fit_setup():
+    dm = BoringDataModule()
+    dm.setup(stage="fit")
+
+    expected_output = (
+        f"name=random_full, size=256{os.linesep}" f"name=random_train, size=64{os.linesep}" f"name=random_val, size=64"
+    )
+    output = str(dm)
+
+    assert expected_output == output
+
+
+def test_datamodule_string_validation_setup():
+    dm = BoringDataModule()
+    dm.setup(stage="validate")
+
+    expected_output = f"name=random_full, size=256{os.linesep}" f"name=random_val, size=64"
+    output = str(dm)
+
+    assert expected_output == output
+
+
+def test_datamodule_string_test_setup():
+    dm = BoringDataModule()
+    dm.setup(stage="test")
+
+    expected_output = f"name=random_full, size=256{os.linesep}" f"name=random_test, size=64"
+    output = str(dm)
+
+    assert expected_output == output
+
+
+def test_datamodule_string_predict_setup():
+    dm = BoringDataModule()
+    dm.setup(stage="predict")
+
+    expected_output = f"name=random_full, size=256{os.linesep}" f"name=random_predict, size=64"
+    output = str(dm)
+
+    assert expected_output == output