Is it possible to extend the ` op` decorator to introduce so

Question

Is it possible to extend the ` op` decorator to introduce some common setup code I don t think a resource can do it since I want access to the op s `context` and a single process executor wouldn t call it for each op I tried making a wrapping decorator but I get errors from Dagster s inspection of the argument lists ```def my op args kwargs Dagster op that opens a custom context manager based on some shared state def wrapper func op args kwargs def raw op context args kwrags my wrapper that does a little extra work with my custom setup context func context args kwargs the original function that would normally be decorated with op return raw op return wrapper```

alex · Answer

gt I get errors from Dagster s inspection of the argument lists if you use `wraps` `update wrapper` it may get the function metadata corrected <https docs python org 3 library functools html functools wraps>

Mark Fickett · Answer

That worked great for posterity I ended up with ```def my op decorator args decorator kwargs def wrapper func op decorator args decorator kwargs wraps func If the wrapped op is missing a context argument this will get an error like raw op missing 1 required positional argument context But we need the context below and want Dagster to pass args according to the wrapped functions args so just require all otel ops to take context def raw op context func args func kwargs with do my custom setup context return func context func args func kwargs return raw op return wrapper```

Mark Fickett · Answer

This is mostly working but when I have a ` my op` that yields multiple outputs I m having trouble I wanted the decorated `func` to be inside the `do my custom setup` context manager That works fine with ` my op`s with a single return value if I look at the stack I see `raw op` in the call stack But when I `yield` if I look at the stack within my decorated `func` I see iterator management and a `spawn` but I don t see `raw op` anywhere in my stack Is there a way I can adjust this so that can have `raw op` in my call stack for the decorated `func` This works single return ``` gt for event in iterate with context home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster utils init py 401 iterate with context gt next output = next iterator home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core execution plan compute generator py 65 coerce solid compute fn to iterator gt result = fn context kwargs if context arg provided else fn kwargs home mfickett Documents mfickett dev data pipeline orchestration op util py 128 raw op lt where I set my context manager gt return func context func args func kwargs gt home mfickett Documents mfickett dev data pipeline orchestration main py 163 list test ids by pipe lt my decorated function``` But this doesn t work ``` Pdb w lt string gt 1 lt module gt usr lib python3 9 multiprocessing spawn py 116 spawn main gt exitcode = main fd parent sentinel usr lib python3 9 multiprocessing spawn py 129 main gt return self bootstrap parent sentinel usr lib python3 9 multiprocessing process py 315 bootstrap gt self run usr lib python3 9 multiprocessing process py 108 run gt self target self args self kwargs home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core executor child process executor py 70 execute command in child process gt for step event in command execute home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core executor multiprocess py 82 execute gt yield from execute plan iterator home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core execution api py 883 iter gt yield from self iterator home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core execution plan execute plan py 87 inner plan execution iterator gt for step event in check generator dagster event sequence for step step context home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core execution plan execute plan py 232 dagster event sequence for step gt for step event in check generator step events home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core execution plan execute step py 354 core dagster event sequence for step gt for user event in check generator home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core execution plan execute step py 70 step output error checked user event sequence gt for user event in user event sequence home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core execution plan compute py 170 execute core compute gt for step output in yield compute results step context inputs compute fn home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core execution plan compute py 138 yield compute results gt for event in iterate with context home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster utils init py 401 iterate with context gt next output = next iterator home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core execution plan compute generator py 66 coerce solid compute fn to iterator gt for event in validate and coerce solid result to iterator result context output defs home mfickett Documents mfickett dev data pipeline direnv python 3 9 5 lib python3 9 site packages dagster core execution plan compute generator py 86 validate and coerce solid result to iterator gt for event in result gt home mfickett Documents mfickett dev data pipeline orchestration main py 264 decide whether to do metadata lt my decorated function but no raw op in its stack```

alex · Answer

context managers and iterators are tricky business What is happening is in the `return func ` line you are returning an iterator generator object and closing the context manager since return exits that scope you will have to rewrite things such that you consume the iterator within the scope of the context manager approximately untested ``` def raw op context func args func kwargs with do my custom setup context res = func context func args func kwargs if inspect isgenerator res for x in res yield x else return res return raw op```

Mark Fickett · Answer

Ah thanks That s much less mysterious than I was worried it could be

Mark Fickett · Answer

I just had to change `return res` to `yield Output res result ` ok hand skin tone 2