def prepare_data(pks):
"""
Return the task instance, data model path, and spectrum for each given primary key,
and apply any spectrum callbacks to the spectrum as it is loaded.
:param pks:
Primary keys of task instances to load data products for.
:returns:
Yields a four length tuple containing the task instance, the spectrum path, the
original spectrum, and the modified spectrum after any spectrum callbacks have been
executed. If no spectrum callback is executed, then the modified spectrum will be
`None`.
"""
trees = {}
for pk in deserialize_pks(pks, flatten=True):
q = session.query(
astradb.TaskInstance).filter(astradb.TaskInstance.pk == pk)
instance = q.one_or_none()
if instance is None:
log.warning(f"No task instance found for primary key {pk}")
path = spectrum = None
else:
release = instance.parameters["release"]
tree = trees.get(release, None)
if tree is None:
trees[release] = tree = SDSSPath(release=release)
# Monkey-patch BOSS Spec paths.
try:
path = tree.full(**instance.parameters)
except:
if instance.parameters["filetype"] == "spec":
from astra.utils import monkey_patch_get_boss_spec_path
path = monkey_patch_get_boss_spec_path(
**instance.parameters)
else:
raise
try:
spectrum = Spectrum1D.read(path)
except:
log.exception(
f"Unable to load Spectrum1D from path {path} on task instance {instance}"
)
spectrum = None
else:
# Are there any spectrum callbacks?
spectrum_callback = instance.parameters.get(
"spectrum_callback", None)
if spectrum_callback is not None:
spectrum_callback_kwargs = instance.parameters.get(
"spectrum_callback_kwargs", "{}")
try:
spectrum_callback_kwargs = literal_eval(
spectrum_callback_kwargs)
except:
log.exception(
f"Unable to literally evalute spectrum callback kwargs for {instance}: {spectrum_callback_kwargs}"
)
raise
try:
func = string_to_callable(spectrum_callback)
spectrum = func(spectrum=spectrum,
path=path,
instance=instance,
**spectrum_callback_kwargs)
except:
log.exception(
f"Unable to execute spectrum callback '{spectrum_callback}' on {instance}"
)
raise
yield (instance, path, spectrum)