diff --git a/metaflow/tests/flows/step_counter_flow.py b/metaflow/tests/flows/step_counter_flow.py new file mode 100644 index 00000000000..5ad1cfbb6e3 --- /dev/null +++ b/metaflow/tests/flows/step_counter_flow.py @@ -0,0 +1,18 @@ +from metaflow.plugins import STEP_DECORATORS + +from metaflow.tests.flows import NewLinearFlow +from metaflow.tests.step_counter import StepCounter + +STEP_DECORATORS.append(StepCounter) + + +class StepCounterFlow(NewLinearFlow): + """Example flow that can receive `--with step_counter` from the CLI + + In order for `--with my_decorator` to work on the CLI, `my_decorator` needs to added to `STEP_DECORATORS`, as above. + I don't know of a good way to accomplish that other than making a trivial flow (`StepCounterFlow` here) that just + mixes in the target flow (`NewLinearFlow`) in a file that appends the decorator (`StepCounter`) to `STEP_DECORATORS. + + TODO: formalize an API for registering step- and flow-decorators without having to define a new flow like this. + """ + pass diff --git a/metaflow/tests/step_counter.py b/metaflow/tests/step_counter.py new file mode 100644 index 00000000000..91a3ee6bf6c --- /dev/null +++ b/metaflow/tests/step_counter.py @@ -0,0 +1,10 @@ +from metaflow.decorators import StepDecorator + +class StepCounter(StepDecorator): + name = "step_counter" + + def task_post_step(self, step_name, flow, graph, retry_count, max_user_code_retries): + if not hasattr(flow, 'step_count'): + flow.step_count = 0 + flow.step_count += 1 + print('step count: %d' % flow.step_count) diff --git a/metaflow/tests/test_decorators.py b/metaflow/tests/test_decorators.py new file mode 100644 index 00000000000..315c5a3397c --- /dev/null +++ b/metaflow/tests/test_decorators.py @@ -0,0 +1,8 @@ +from metaflow.tests.flows.step_counter_flow import StepCounterFlow +from metaflow.tests.utils import run + + +def test_step_count(): + flow = StepCounterFlow + data = run(flow, args=('run','--with','step_counter',)) + assert data == { 'a': 111, 'b': 222, 'checked': True, 'step_count': 5, }