def test_register_operation():
"""
The most simple way to extend the pipeline operations is by passing a
function with proper arguments.
:return:
"""
test_pipe = Pipeline(STEPS, models=[('ents', 'nl', model_path_nl), ('other_identifier', 'en', model_path_en)])
def custom_op(doc, **kwargs):
return kwargs
test_pipe.register_operation('CUSTOM_STEP', custom_op)
assert len(test_pipe._operations) == len(STEPS) + 1
def test_register_op_with_extending_steps_works():
"""
Calling the custom pipeline operation with an argument should yield the same
arguments passed back as a result
:return:
"""
test_pipe = Pipeline(STEPS, **PIPELINE_DEF_KWARGS)
def custom_op(doc, context=None, settings=None, **kwargs):
return settings
custom_argument = {'argument': 1}
test_pipe.register_operation('CUSTOM_STEP', custom_op)
test_pipe.steps.append(('CUSTOM_STEP', custom_argument))
results = test_pipe(TEXT)
assert results['CUSTOM_STEP'] == custom_argument
def test_context_data_passed_between_steps():
"""
Calling the custom pipeline operation with an argument should yield the same
arguments passed back as a result
:return:
"""
test_pipe = Pipeline([], **PIPELINE_DEF_KWARGS)
def custom_op(doc, context=None, settings=None, **kwargs):
return 1
def custom_op2(doc, context=None, settings=None, **kwargs):
return context
custom_argument = {'argument': 1}
test_pipe.register_operation('CUSTOM_STEP', custom_op)
test_pipe.register_operation('CUSTOM_STEP2', custom_op2)
test_pipe.steps.append(('CUSTOM_STEP', custom_argument))
test_pipe.steps.append(('CUSTOM_STEP2', custom_argument))
results = test_pipe(TEXT)
assert results['CUSTOM_STEP2']['CUSTOM_STEP'] == results['CUSTOM_STEP']