I've used Python's context managers frequently, usually in place of some try-catch-finally type of logic. However, recently I've seen some examples of these where they go beyond the simple, and I'm trying to understand how they work. As an example, here's a pattern I've seen in some libraries
first = SomeClass()
second = SomeOtherClass()
with SomeContext() as c:
first.do_calculation()
second.do_calculation()
c.execute()
In these examples the context manager is somehow aware of the lines of code it contextualises without explicitly passing the variable 'c' to the methods. A concrete example can be seen in for example the Prefect workflow engine where tasks are added to a flow which is then executed. Here's a simple example
from prefect import task, Flow
@task
def first(x, y):
return x + y
@task
def second(x, y):
return x*y
with Flow('test-flow') as flow:
res_first = first(1, 2)
res_second = second(res_first)
flow.run()
Here somehow the flow 'test-flow' becomes aware of the tasks within its context, but I really don't understand how. The example I'm most closely looking to mirror is that of gs-quant which has this type of logic for pricing financial derivatives
ir_swap = InterestRateSwap(...) # fictive class names for brevity
eq_option = EquityOption(...)
with PricingContext(...) as p:
ir_price = ir_swap.price()
eq_price = eq_option.price()
print(ir_price.result(), eq_price.result())
Here the PricingContext could e.g. set market rates for another date, thus altering the return of the .price() call. The guides for this library mention the following:
Note that using a PricingContext as a context manager has two extra effects: All calls to price(), calc() are dispatched as a single request, on context manager exit. This allows for the communication overhead to be borne only once for multiple calculations.
So, as an example, suppose I have the following code
class MyClass:
def __init__(self):
...
def calculate(self):
...
my_class = MyClass()
with MyContext() as c:
my_class.calculate()
where I want the context to inform the logic of .calculate()
I'm aware one could simply pass the context manager as an argument to the method, but the libraries I mentioned don't seem to do that. I'm assuming the libraries mentioned above have some sort of (global) default manager set but I'm not sure what a good design for that would be.
I have tried to understand the logic that both prefect as gs-quant use, however their codebases are relatively dense and difficult to parse (lots of metaclasses and such)
As @chepner wrote, there is no straightforward (or implicit) syntax to make this work in python, at least not at the moment. @barmar's comment was on-the-money (the docstring for def task(): Decorator to designate a function as a task in a Prefect workflow.), so if gs-quant's code is a bit dense, maybe Prefect's definition for task can make for an easier read (I didn't check the former myself).
Coming back to @chepner's second comment, it is not impossible to make the context manager be aware of underlying executions, but that would usually be much more complicated to achieve, and significantly less maintainable. So as they wrote - for most implementations (and as a general pattern), it is probably more accurate to say that the code in the body is aware of a setting/state which is set by the running context, not the other way around. And with this, I'll proceed to the first example of how body-code can cooperate with the context it is running in;
The following file (let's call it mul_simple.py) shows how a context can "affect" underlying code. More accurately put, it shows how underlying code can interact with the context it is running in. In this example, print_result receives a textual representation, and will use the context to multiply num by its setting to match the text. The context manager tracks its current state (to allow context re-entry with different values), and exposes its helper function for the use of underlying code. While the context performs the multiplication itself, it can simply expose the current multiplier to a calling function for use at its own discretion.
def print_result(num_name, func, num):
result = func(num).calculate()
print(f"{num_name} is: {result}")
class Multiplier:
_multiply_context = [1]
@classmethod
def get_result(cls, multiplied):
return multiplied * cls._multiply_context[-1]
def __init__(self, multiplier):
self.multiplier = int(multiplier)
def __enter__(self):
self._multiply_context.append(self.multiplier)
def __exit__(self, type_, value_, traceback_):
_ = self._multiply_context.pop()
if any(arg is not None for arg in (type_, value_, traceback_)):
return False # reraise
class MultipliedBase:
def calculate(self):
return Multiplier.get_result(self)
class MultipliedInt(int, MultipliedBase):
pass
def main_int():
func = MultipliedInt
print_result("Three", func, 3)
with Multiplier(2):
print_result("Six", func, 3)
with Multiplier(4):
print_result("Sixteen", func, 4)
print_result("Eight", func, 4)
print_result("Five", func, 5)
if __name__ == "__main__":
main_int()
Running this will output:
$ python mul_simple.py
Three is: 3
Six is: 6
Sixteen is: 16
Eight is: 8
Five is: 5
Say we wanted to extend our context to multiply strings as well. This is one case where had we made the context aware of the body-code - we'd have to make more significant, delicate changes, to the context manager itself. Luckily we chose the other path, so no changes need be performed, only adding a new file. Let's call it mul_str.py:
from mul_simple import print_result, Multiplier, MultipliedBase
class MultipliedStr(str, MultipliedBase):
pass
def main_str():
func = MultipliedStr
print_result("One", func, "One")
with Multiplier(2):
print_result("AbAb", func, "Ab")
with Multiplier(4):
print_result("FourFourFourFour", func, "Four")
print_result("FooFoo", func, "Foo")
print_result("Bar", func, "Bar")
if __name__ == "__main__":
main_str()
And running it will result in:
$ python mul_str.py
One is: One
AbAb is: AbAb
FourFourFourFour is: FourFourFourFour
FooFoo is: FooFoo
Bar is: Bar
To demonstrate how this can work as a function decorator (like @task in Prefect), we can create a file called mul_deco.py, as an example:
from mul_simple import Multiplier
class MultiplierWithDecorator(Multiplier):
@classmethod
def multiply(cls, multiplier):
def outer(func):
def inner(*args, **kwargs):
with cls(multiplier):
return cls.get_result(func(*args, **kwargs))
return inner
return outer
@MultiplierWithDecorator.multiply(3)
def by_six(arg): # arg * 2 * 3
return arg * 2
if __name__ == "__main__":
print(f"Thirty six is: {by_six(6)}")
print(f"Forty two is: {by_six(7)}")
And executing it will print:
$ python mul_deco.py
Thirty six is: 36
Forty two is: 42
Finally, while the context manager won't be aware of the underlying code, we can make it collect multiplication request for us to execute at a later time. With slight adjustments this can be utilized to aggregate requests over a single connection. If the contents of the file mul_aggregate.py are as such:
class MultiplyLater:
_multiply_context = []
@classmethod
def add_result(cls, name, to_multiply):
try:
cls._multiply_context[-1]._delayed_results.append((name, to_multiply))
except IndexError:
raise IndexError("No active multiplication context")
def __init__(self, multiplier):
self.multiplier = int(multiplier)
self._delayed_results = []
def __enter__(self):
self._multiply_context.append(self)
return self
def __exit__(self, type_, value_, traceback_):
_ = self._multiply_context.pop()
if any(arg is not None for arg in (type_, value_, traceback_)):
return False # reraise
def get_results(self):
return [(name, to_multiply * self.multiplier) for (name, to_multiply) in self._delayed_results]
class LaterMultiplied:
def __init__(self, name, to_multiply):
MultiplyLater.add_result(name, to_multiply)
if __name__ == "__main__":
with MultiplyLater(3) as later_three:
LaterMultiplied("Six", 2)
LaterMultiplied("Nine", 3)
with MultiplyLater(2) as later_two:
LaterMultiplied("Four", 2)
LaterMultiplied("Twelve", 6)
LaterMultiplied("Fifteen", 5)
for name, result in later_three.get_results():
print(f"By 3; {name} is: {result}")
for name, result in later_two.get_results():
print(f"By 2; {name} is: {result}")
Then running it will print:
$ python mul_aggregate.py
By 3; Six is: 6
By 3; Nine is: 9
By 3; Fifteen is: 15
By 2; Four is: 4
By 2; Twelve is: 12
While this is still far from how libraries like Prefect implement their capabilities, I'm hoping this demo provides a foundation which makes some sense.