Exposure changes and creation of ndsl.typing module

ndsl/__init__.py

@@ -1,29 +1,10 @@
 from .buffer import Buffer
-from .checkpointer.base import Checkpointer
-from .checkpointer.null import NullCheckpointer
-from .checkpointer.snapshots import SnapshotCheckpointer, _Snapshots
-from .checkpointer.thresholds import (
-    InsufficientTrialsError,
-    SavepointThresholds,
-    Threshold,
-    ThresholdCalibrationCheckpointer,
-)
-from .checkpointer.validation import ValidationCheckpointer
 from .comm.boundary import Boundary, SimpleBoundary
-from .comm.caching_comm import (
-    CachingCommData,
-    CachingCommReader,
-    CachingCommWriter,
-    CachingRequestReader,
-    CachingRequestWriter,
-    NullRequest,
-)
-from .comm.comm_abc import Comm, Request
-from .comm.communicator import Communicator, CubedSphereCommunicator, TileCommunicator
+from .comm.communicator import CubedSphereCommunicator, TileCommunicator
 from .comm.local_comm import AsyncResult, ConcurrencyError, LocalComm
 from .comm.mpi import MPIComm
 from .comm.null_comm import NullAsyncResult, NullComm
-from .comm.partitioner import CubedSpherePartitioner, Partitioner, TilePartitioner
+from .comm.partitioner import CubedSpherePartitioner, TilePartitioner
 from .constants import ConstantVersions
 from .dsl.caches.codepath import FV3CodePath
 from .dsl.dace.dace_config import DaceConfig, DaCeOrchestration, FrozenCompiledSDFG
@@ -44,17 +25,6 @@
 )
 from .dsl.stencil_config import CompilationConfig, RunMode, StencilConfig
 from .exceptions import OutOfBoundsError
-from .grid.eta import HybridPressureCoefficients
-from .grid.generation import GridDefinition, GridDefinitions, MetricTerms
-from .grid.helper import (
-    AngleGridData,
-    ContravariantGridData,
-    DampingCoefficients,
-    DriverGridData,
-    GridData,
-    HorizontalGridData,
-    VerticalGridData,
-)
 from .halo.data_transformer import (
     HaloDataTransformer,
     HaloDataTransformerCPU,
@@ -86,24 +56,6 @@
     QuantityHaloSpec,
     QuantityMetadata,
 )
-from .stencils.c2l_ord import CubedToLatLon
-from .stencils.corners import CopyCorners, CopyCornersXY, FillCornersBGrid
-from .stencils.testing.grid import Grid  # type: ignore
-from .stencils.testing.parallel_translate import (
-    ParallelTranslate,
-    ParallelTranslate2Py,
-    ParallelTranslate2PyState,
-    ParallelTranslateBaseSlicing,
-    ParallelTranslateGrid,
-)
-from .stencils.testing.savepoint import SavepointCase, Translate, dataset_to_dict
-from .stencils.testing.temporaries import assert_same_temporaries, copy_temporaries
-from .stencils.testing.translate import (
-    TranslateFortranData2Py,
-    TranslateGrid,
-    pad_field_in_j,
-    read_serialized_data,
-)
 from .testing.dummy_comm import DummyComm
 from .types import Allocator, AsyncRequest, NumpyModule
 from .units import UnitsError

ndsl/checkpointer/__init__.py

@@ -0,0 +1,9 @@
+from .null import NullCheckpointer
+from .snapshots import SnapshotCheckpointer, _Snapshots
+from .thresholds import (
+    InsufficientTrialsError,
+    SavepointThresholds,
+    Threshold,
+    ThresholdCalibrationCheckpointer,
+)
+from .validation import ValidationCheckpointer

ndsl/checkpointer/null.py

@@ -1,4 +1,4 @@
-from ndsl.checkpointer.base import Checkpointer
+from ndsl.typing import Checkpointer
 
 
 class NullCheckpointer(Checkpointer):

ndsl/checkpointer/snapshots.py

@@ -2,9 +2,9 @@
 
 import numpy as np
 
-from ndsl.checkpointer.base import Checkpointer
 from ndsl.optional_imports import cupy as cp
 from ndsl.optional_imports import xarray as xr
+from ndsl.typing import Checkpointer
 
 
 def make_dims(savepoint_dim, label, data_list):

ndsl/checkpointer/thresholds.py

@@ -5,8 +5,8 @@
 
 import numpy as np
 
-from ndsl.checkpointer.base import Checkpointer
 from ndsl.quantity import Quantity
+from ndsl.typing import Checkpointer
 
 
 try:

ndsl/checkpointer/validation.py

@@ -5,14 +5,14 @@
 
 import numpy as np
 
-from ndsl.checkpointer.base import Checkpointer
 from ndsl.checkpointer.thresholds import (
     ArrayLike,
     SavepointName,
     SavepointThresholds,
     cast_to_ndarray,
 )
 from ndsl.optional_imports import xarray as xr
+from ndsl.typing import Checkpointer
 
 
 def _clip_pace_array_to_target(

ndsl/comm/__init__.py

@@ -0,0 +1,9 @@
+from .caching_comm import (
+    CachingCommData,
+    CachingCommReader,
+    CachingCommWriter,
+    CachingRequestReader,
+    CachingRequestWriter,
+    NullRequest,
+)
+from .comm_abc import Comm, Request