From 1b72bafa98c0146f08d89e67bf8c8327ef80b45b Mon Sep 17 00:00:00 2001
From: Adam Narozniak <51029327+adam-narozniak@users.noreply.github.com>
Date: Thu, 20 Jun 2024 12:27:15 +0200
Subject: [PATCH] fix(datasets) Fix sorting in `__all__` (#3655)

---
 datasets/flwr_datasets/__init__.py                   |  5 +++--
 datasets/flwr_datasets/partitioner/__init__.py       | 12 ++++++------
 datasets/flwr_datasets/preprocessor/__init__.py      |  2 +-
 datasets/flwr_datasets/visualization/__init__.py     |  2 +-
 datasets/flwr_datasets/visualization/bar_plot.py     |  1 -
 datasets/flwr_datasets/visualization/heatmap_plot.py |  2 --
 6 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/datasets/flwr_datasets/__init__.py b/datasets/flwr_datasets/__init__.py
index d084780102ce..bd68fa43c606 100644
--- a/datasets/flwr_datasets/__init__.py
+++ b/datasets/flwr_datasets/__init__.py
@@ -23,11 +23,12 @@
 
 __all__ = [
     "FederatedDataset",
-    "partitioner",
     "metrics",
-    "visualization",
+    "partitioner",
     "preprocessor",
     "utils",
+    "visualization",
 ]
 
+
 __version__ = _package_version
diff --git a/datasets/flwr_datasets/partitioner/__init__.py b/datasets/flwr_datasets/partitioner/__init__.py
index 0d1edbfcb04a..1fc00ed90323 100644
--- a/datasets/flwr_datasets/partitioner/__init__.py
+++ b/datasets/flwr_datasets/partitioner/__init__.py
@@ -27,14 +27,14 @@
 from .square_partitioner import SquarePartitioner
 
 __all__ = [
+    "DirichletPartitioner",
+    "ExponentialPartitioner",
     "IidPartitioner",
-    "Partitioner",
+    "InnerDirichletPartitioner",
+    "LinearPartitioner",
     "NaturalIdPartitioner",
-    "DirichletPartitioner",
+    "Partitioner",
+    "ShardPartitioner",
     "SizePartitioner",
-    "LinearPartitioner",
-    "InnerDirichletPartitioner",
     "SquarePartitioner",
-    "ShardPartitioner",
-    "ExponentialPartitioner",
 ]
diff --git a/datasets/flwr_datasets/preprocessor/__init__.py b/datasets/flwr_datasets/preprocessor/__init__.py
index bab5d82a2035..67b2aaadc3d2 100644
--- a/datasets/flwr_datasets/preprocessor/__init__.py
+++ b/datasets/flwr_datasets/preprocessor/__init__.py
@@ -20,7 +20,7 @@
 from .preprocessor import Preprocessor
 
 __all__ = [
+    "Divider",
     "Merger",
     "Preprocessor",
-    "Divider",
 ]
diff --git a/datasets/flwr_datasets/visualization/__init__.py b/datasets/flwr_datasets/visualization/__init__.py
index 801a38dcafc6..b55e406c71db 100644
--- a/datasets/flwr_datasets/visualization/__init__.py
+++ b/datasets/flwr_datasets/visualization/__init__.py
@@ -19,6 +19,6 @@
 from .label_distribution import plot_label_distributions
 
 __all__ = [
-    "plot_label_distributions",
     "plot_comparison_label_distribution",
+    "plot_label_distributions",
 ]
diff --git a/datasets/flwr_datasets/visualization/bar_plot.py b/datasets/flwr_datasets/visualization/bar_plot.py
index 6326b24a9695..339ff0967906 100644
--- a/datasets/flwr_datasets/visualization/bar_plot.py
+++ b/datasets/flwr_datasets/visualization/bar_plot.py
@@ -38,7 +38,6 @@ def _plot_bar(
     plot_kwargs: Optional[Dict[str, Any]],
     legend_kwargs: Optional[Dict[str, Any]],
 ) -> Axes:
-
     if axis is None:
         if figsize is None:
             figsize = _initialize_figsize(
diff --git a/datasets/flwr_datasets/visualization/heatmap_plot.py b/datasets/flwr_datasets/visualization/heatmap_plot.py
index 2e593a79368e..3c87de7693ae 100644
--- a/datasets/flwr_datasets/visualization/heatmap_plot.py
+++ b/datasets/flwr_datasets/visualization/heatmap_plot.py
@@ -39,7 +39,6 @@ def _plot_heatmap(
     plot_kwargs: Optional[Dict[str, Any]],
     legend_kwargs: Optional[Dict[str, Any]],
 ) -> Axes:
-
     if axis is None:
         if figsize is None:
             figsize = _initialize_figsize(
@@ -92,7 +91,6 @@ def _initialize_figsize(
     num_partitions: int,
     num_labels: int,
 ) -> Tuple[float, float]:
-
     figsize = (0.0, 0.0)
     if partition_id_axis == "x":
         figsize = (3 * np.sqrt(num_partitions), np.sqrt(num_labels))