Be explicit about internal vs external API

dask · Oct 25, 2022 · 400cc5a · 400cc5a
1 parent 05121ca
commit 400cc5a
Show file tree

Hide file tree

Showing 11 changed files with 17 additions and 12 deletions.
diff --git a/distributed/shuffle/__init__.py b/distributed/shuffle/__init__.py
@@ -1,8 +1,13 @@
 from __future__ import annotations
 
-from distributed.shuffle.shuffle import rearrange_by_column_p2p
-from distributed.shuffle.shuffle_extension import (
-    ShuffleId,
+from distributed.shuffle._shuffle import rearrange_by_column_p2p
+from distributed.shuffle._shuffle_extension import (
     ShuffleSchedulerExtension,
     ShuffleWorkerExtension,
 )
+
+__all__ = [
+    "ShuffleSchedulerExtension",
+    "ShuffleWorkerExtension",
+    "rearrange_by_column_p2p",
+]
diff --git a/distributed/shuffle/arrow.py → distributed/shuffle/_arrow.py b/distributed/shuffle/arrow.py → distributed/shuffle/_arrow.py
diff --git a/distributed/shuffle/multi_comm.py → distributed/shuffle/_multi_comm.py b/distributed/shuffle/multi_comm.py → distributed/shuffle/_multi_comm.py
diff --git a/distributed/shuffle/multi_file.py → distributed/shuffle/_multi_file.py b/distributed/shuffle/multi_file.py → distributed/shuffle/_multi_file.py
diff --git a/distributed/shuffle/shuffle.py → distributed/shuffle/_shuffle.py b/distributed/shuffle/shuffle.py → distributed/shuffle/_shuffle.py
@@ -6,7 +6,7 @@
 from dask.delayed import Delayed
 from dask.highlevelgraph import HighLevelGraph, MaterializedLayer
 
-from distributed.shuffle.shuffle_extension import ShuffleId, ShuffleWorkerExtension
+from distributed.shuffle._shuffle_extension import ShuffleId, ShuffleWorkerExtension
 
 if TYPE_CHECKING:
     import pandas as pd

diff --git a/distributed/shuffle/shuffle_extension.py → distributed/shuffle/_shuffle_extension.py b/distributed/shuffle/shuffle_extension.py → distributed/shuffle/_shuffle_extension.py
@@ -22,14 +22,14 @@
 import toolz
 
 from distributed.protocol import to_serialize
-from distributed.shuffle.arrow import (
+from distributed.shuffle._arrow import (
     deserialize_schema,
     dump_batch,
     list_of_buffers_to_table,
     load_arrow,
 )
-from distributed.shuffle.multi_comm import MultiComm
-from distributed.shuffle.multi_file import MultiFile
+from distributed.shuffle._multi_comm import MultiComm
+from distributed.shuffle._multi_file import MultiFile
 from distributed.utils import log_errors, sync
 
 if TYPE_CHECKING:

diff --git a/distributed/shuffle/tests/test_graph.py b/distributed/shuffle/tests/test_graph.py
@@ -13,7 +13,7 @@
 from dask.blockwise import Blockwise
 from dask.utils_test import hlg_layer_topological
 
-from distributed.shuffle.shuffle_extension import ShuffleWorkerExtension
+from distributed.shuffle._shuffle_extension import ShuffleWorkerExtension
 from distributed.utils_test import gen_cluster
 
 

diff --git a/distributed/shuffle/tests/test_multi_comm.py b/distributed/shuffle/tests/test_multi_comm.py
@@ -5,7 +5,7 @@
 
 import pytest
 
-from distributed.shuffle.multi_comm import MultiComm
+from distributed.shuffle._multi_comm import MultiComm
 from distributed.utils_test import gen_test
 
 

diff --git a/distributed/shuffle/tests/test_multi_file.py b/distributed/shuffle/tests/test_multi_file.py
@@ -5,7 +5,7 @@
 
 import pytest
 
-from distributed.shuffle.multi_file import MultiFile
+from distributed.shuffle._multi_file import MultiFile
 from distributed.utils_test import gen_test
 
 

diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
@@ -17,7 +17,7 @@
 from dask.distributed import Worker
 from dask.utils import stringify
 
-from distributed.shuffle.shuffle_extension import (
+from distributed.shuffle._shuffle_extension import (
     dump_batch,
     list_of_buffers_to_table,
     load_arrow,

diff --git a/distributed/shuffle/tests/test_shuffle_extension.py b/distributed/shuffle/tests/test_shuffle_extension.py
@@ -5,7 +5,7 @@
 pd = pytest.importorskip("pandas")
 dd = pytest.importorskip("dask.dataframe")
 
-from distributed.shuffle.shuffle_extension import (
+from distributed.shuffle._shuffle_extension import (
     ShuffleWorkerExtension,
     get_worker_for,
     split_by_partition,