Add feature implementation to datamodule for str method

First implementation scetch

Author: MrWhatZitToYaa

src/lightning/pytorch/core/datamodule.py

@@ -16,28 +16,27 @@
 import inspect
 from typing import IO, Any, Dict, Iterable, Optional, Union, cast
 
+import pytorch_lightning as pl
+from lightning_fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
 from lightning_utilities import apply_to_collection
+from pytorch_lightning.core.hooks import DataHooks
+from pytorch_lightning.core.mixins import HyperparametersMixin
+from pytorch_lightning.core.saving import _load_from_checkpoint
+from pytorch_lightning.utilities.model_helpers import _restricted_classmethod
+from pytorch_lightning.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
 from torch.utils.data import DataLoader, Dataset, IterableDataset
 from typing_extensions import Self
 
-import lightning.pytorch as pl
-from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
-from lightning.pytorch.core.hooks import DataHooks
-from lightning.pytorch.core.mixins import HyperparametersMixin
-from lightning.pytorch.core.saving import _load_from_checkpoint
-from lightning.pytorch.utilities.model_helpers import _restricted_classmethod
-from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
-
 
 class LightningDataModule(DataHooks, HyperparametersMixin):
     """A DataModule standardizes the training, val, test splits, data preparation and transforms. The main advantage is
     consistent data splits, data preparation and transforms across models.
 
     Example::
 
-        import lightning as L
+        import lightning.pytorch as L
         import torch.utils.data as data
-        from lightning.pytorch.demos.boring_classes import RandomDataset
+        from pytorch_lightning.demos.boring_classes import RandomDataset
 
         class MyDataModule(L.LightningDataModule):
             def prepare_data(self):
@@ -243,3 +242,32 @@ 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 = []
+
+        for attr_name in dir(self):
+            attr = getattr(self, attr_name)
+
+            # Get Dataset information
+            if isinstance(attr, Dataset):
+                if hasattr(attr, "__len__"):
+                    datasets_info.append(f"{attr_name}, dataset size={len(attr)}")
+                else:
+                    datasets_info.append(f"{attr_name}, dataset size=Unavailable")
+            elif isinstance(attr, (list, tuple)) and all(isinstance(item, Dataset) for item in attr):
+                if all(hasattr(item, "__len__") for item in attr):
+                    datasets_info.append(f"{attr_name}, dataset size={[len(ds) for ds in attr]}")
+                else:
+                    datasets_info.append(f"{attr_name}, dataset size=Unavailable")
+
+        if not datasets_info:
+            return "No datasets are set up."
+
+        return "\n".join(datasets_info)