I know this is not necessarily a problem question but more of a style question. I understand that it doesn't have to be pretty for the most part.
In general, context managers are used for mutexing, acquiring resources, and many other things. The context managers usually yield the acquired resource or just provide wrapping functionality.
In general, generators are very capable us directing flow with conditionally yielding and custom iteration.
In my situation, I am writing a pipeline for an automatic probe station with a source and sink, but the probe station may or may not be able to probe multiple pins at the same time. That is, either iterating through each pair or being capable of connecting all pairs at the same time.
For a given pin group there is a list or pin pairs to be stressed. There is also a set of corresponding pins that are tested to determine pass or fail of the pin group given the stress level.
I have implemented a context manager for connecting to pins, yielding for testing, and then disconnecting on clean up.
Given a flag for the moment, I can control whether the loop connects all at once versus iterating through each. In the actual implementation the probe_context_generator will live in a concrete class that will offer only what is needed. The multi flag would only be offered if the concrete class can actually do either
Here is the code.
import contextlib
import random
import time
from itertools import product
from typing import ContextManager, Generator
# Business Logic
@contextlib.contextmanager
def probe_context(p): # Belongs to device specific probe station class
try:
print(f'Connecting pins to {p}')
time.sleep(0.1)
print(f'Connected')
yield p
finally:
print(f'Disconnecting {p}')
time.sleep(0.1)
print('Disconnected')
def probe_context_generator(*pin_pairs, multi=False):
if multi:
yield probe_context(pin_pairs)
else:
for pin_pair in pin_pairs:
yield probe_context(pin_pair)
# Application Code
def probe_for_pulse(multi, p_level, pulse_pin_pairs):
for probe_pulse_pins_con in probe_context_generator(*pulse_pin_pairs, multi=multi):
with probe_pulse_pins_con as pulse_pins:
yield p_level, pulse_pins
def check_leakage(multi, leakage_pin_pairs):
leakage_results = []
for probe_leakage_pins_con in probe_context_generator(*leakage_pin_pairs, multi=multi):
with probe_leakage_pins_con as leakage_pin_pair:
print(f'Checking leakage on {leakage_pin_pair=}')
leakage_results.append(random.choice([True, False]))
return leakage_results
def ready_next_pulse_generator(p_levels, pin_groups: list[PinGroup], multi=False):
for pin_group in pin_groups:
print(f"\n!!!!! Pin Group {pin_group.group_num} !!!!!")
for p_level in p_levels:
# Measure Pass fail
leakage_results = check_leakage(multi, pin_group.leakage_pin_pairs)
print("broadcast leakage data to plots and data saving thread")
if not all(leakage_results):
print(f"Pin Group {pin_group.group_num} failed {p_level}")
break
yield from probe_for_pulse(multi, p_level, pin_group.pulse_pin_pairs)
print(f"Pin Group {pin_group.group_num} passed {p_level}")
PINS = ["A1", "B1", "C1"]
pulse_levels = [125, 250, 500]
class PinGroup:
pulse_pin_pairs = [
(force, sense) for force, sense in product(PINS, PINS) if force != sense
]
leakage_pin_pairs = [("Leakage force pin", "Leakage sense pin")]
def __init__(self, group_num):
self.group_num = group_num
print("##### Probe Single #####")
for pulse_level, connected_pin in ready_next_pulse_generator(pulse_levels, [PinGroup(i) for i in range(5)]):
print(f'\tarm scope for {connected_pin}')
time.sleep(0.1)
print(f"\tpulse {connected_pin=}, {pulse_level=}")
time.sleep(0.1)
print(f"\tread data for {connected_pin}")
time.sleep(0.1)
print("\n" * 2)
print("##### Probe Multi #####")
for pulse_level, connected_pins in ready_next_pulse_generator(pulse_levels, [PinGroup(i) for i in range(5)], multi=True):
print(f'\tarm scope for {connected_pins}')
time.sleep(0.1)
print(f"\tpulse {connected_pins=}, {pulse_level=}")
time.sleep(0.1)
print(f"\tread data for {connected_pins}")
time.sleep(0.1)
My main question is, I guess, is the probing_context sufficient need for utilizing the contextmanager functionality? Would making the arming of the scope and reading the data also contextmanager increase the chances of this code being used correctly, or does it seem like I'm overusing a cool thing? Subjective, I know. But I haven't seen a lot of things like this. The probing_context, for a robotic probe station, will handle path planning and doing the actual moving.