Perhaps a custom compose function can help with these use cases? This series has a few examples of composing computation in Python that might be useful.

https://github.com/mpkocher/Functional-Programming-Technique...

And there's a general list of resources here.

https://github.com/sfermigier/awesome-functional-python