Updated change to pass precommit checks. Includes some refactoring for readability.
MODIFY
pangeo_forge_recipes/executors/beam.py
pangeo_forge_recipes/executors/beam.py
import itertools
from dataclasses import dataclass
-from typing import Optional, Iterable, Tuple, Any, List
+from typing import Any, Iterable, List, Tuple, cast
import apache_beam as beam
-from .base import (
- Pipeline,
- PipelineExecutor,
- Config,
- Stage,
- NoArgumentStageFunction,
- SingleArgumentStageFunction
-)
+from .base import Config, NoArgumentStageFunction, Pipeline, PipelineExecutor, Stage
-def _exec_no_arg_stage(
- last: int, *, current: int, function: NoArgumentStageFunction, config: Config
-) -> int:
+def _no_arg_stage(last: int, *, current: int, func: NoArgumentStageFunction, config: Config) -> int:
"""Execute a NoArgumentStageFunction, ensuring execution order."""
- assert (last + 1) == current, f'stages are executing out of order! On step {current!r}.'
+ assert (last + 1) == current, f"stages are executing out of order! On step {current!r}."
- function(config=config)
+ func(config=config)
return current
@@ -37,38 +28,33 @@ class _SingleArgumentStage(beam.PTransform):
stage: Stage
config: Config
- @staticmethod
- def prepare_stage(last: int, *, current: int, mappable: Iterable) -> Iterable[Tuple[int, Any]]:
+ def prepare_stage(self, last: int) -> Iterable[Tuple[int, Any]]:
"""Propagate current stage to Mappables for parallel execution."""
- assert (last + 1) == current, f'stages are executing out of order! On step {current!r}.'
- return zip(itertools.repeat(current), mappable)
+ assert (last + 1) == self.step, f"stages are executing out of order! On step {self.step!r}."
+ return zip(itertools.repeat(self.step), cast(Iterable, self.stage.mappable))
- @staticmethod
- def exec_stage(
- last: int, arg,
- current: int = -1, function: SingleArgumentStageFunction = _no_op, config: Optional[Config] = None
- ) -> int:
+ def exec_stage(self, last: int, arg: Any) -> int:
"""Execute stage function."""
- assert last == current, f'stages are executing out of order! On step {current!r}.'
+ assert last == self.step, f"stages are executing out of order! On step {self.step!r}."
- function(arg, config=config)
+ self.stage.function(arg, config=self.config) # type: ignore
- return current
+ return self.step
- @staticmethod
- def post_validate(last: List[int], *, current: int) -> int:
+ def post_validate(self, last: List[int]) -> int:
"""Propagate step number for downstream stage validation."""
- assert all([it == current for it in last]), f'stages are executing out of order! On step {current!r}.'
- return current
+ in_current_step = all((it == self.step for it in last))
+ assert in_current_step, f"stages are executing out of order! On step {self.step!r}."
+
+ return self.step
def expand(self, pcoll):
return (
pcoll
- | "Prepare" >> beam.FlatMap(self.prepare_stage, current=self.step, mappable=self.stage.mappable)
- | "Execute" >> beam.MapTuple(self.exec_stage, current=self.step, function=self.stage.function,
- config=self.config)
+ | "Prepare" >> beam.FlatMap(self.prepare_stage)
+ | "Execute" >> beam.MapTuple(self.exec_stage)
| beam.combiners.ToList()
- | "Validate" >> beam.Map(self.post_validate, current=self.step)
+ | "Validate" >> beam.Map(self.post_validate)
)
@@ -80,15 +66,14 @@ class BeamPipelineExecutor(PipelineExecutor[beam.PTransform]):
if stage.mappable is not None:
pcoll |= stage.name >> _SingleArgumentStage(step, stage, pipeline.config)
else:
- pcoll |= stage.name >> beam.Map(_exec_no_arg_stage,
- current=step,
- function=stage.function,
- config=pipeline.config)
+ pcoll |= stage.name >> beam.Map(
+ _no_arg_stage, current=step, func=stage.function, config=pipeline.config
+ )
return pcoll
@staticmethod
def execute(plan: beam.PTransform, *args, **kwargs):
- """Execute a plan. All args and kwargs are pass to a `apache_beam.Pipeline`."""
+ """Execute a plan. All args and kwargs are passed to a `apache_beam.Pipeline`."""
with beam.Pipeline(*args, **kwargs) as p:
p | plan

MODIFY
setup.cfg
setup.cfg
@@ -60,7 +60,7 @@ max-line-length = 100
[isort]
known_first_party=pangeo_forge_recipes
-known_third_party=aiohttp,click,dask,fsspec,fsspec_reference_maker,mypy_extensions,numpy,pandas,prefect,pytest,pytest_lazyfixture,setuptools,xarray,yaml,zarr
+known_third_party=aiohttp,apache_beam,click,dask,fsspec,fsspec_reference_maker,mypy_extensions,numpy,pandas,prefect,pytest,pytest_lazyfixture,setuptools,xarray,yaml,zarr
multi_line_output=3
include_trailing_comma=True
force_grid_wrap=0

MODIFY
tests/test_pipelines.py
tests/test_pipelines.py
@@ -6,10 +6,10 @@ import pytest
from pytest_lazyfixture import lazy_fixture
from pangeo_forge_recipes.executors.base import Pipeline, Stage
+from pangeo_forge_recipes.executors.beam import BeamPipelineExecutor
from pangeo_forge_recipes.executors.dask import DaskPipelineExecutor
from pangeo_forge_recipes.executors.function import FunctionPipelineExecutor
from pangeo_forge_recipes.executors.prefect import PrefectPipelineExecutor
-from pangeo_forge_recipes.executors.beam import BeamPipelineExecutor
@pytest.fixture
@@ -52,7 +52,8 @@ def pipeline_with_config(tmpdir_factory):
@pytest.mark.parametrize(
- "Executor", [FunctionPipelineExecutor, PrefectPipelineExecutor, DaskPipelineExecutor, BeamPipelineExecutor],
+ "Executor",
+ [FunctionPipelineExecutor, PrefectPipelineExecutor, DaskPipelineExecutor, BeamPipelineExecutor],
)
@pytest.mark.parametrize(
"pipeline_config_tmpdir",
@@ -70,4 +71,4 @@ def test_pipeline(pipeline_config_tmpdir, Executor):
f"{prefix}func1_b.log",
f"{prefix}func1_3.log",
]:
- assert tmpdir.join(fname).check(file=True), f'File not found in temporary directory: {fname}.'
+ assert tmpdir.join(fname).check(file=True), f"File not found in temp directory: {fname}."

Updated change to pass precommit checks. Includes some refactoring for readability.