"""
Classes for managing processed data.
"""
__all__ = ["MetricStatus", "ReqReport", "Storable"]
import datetime
import importlib
from abc import abstractmethod
from dataclasses import dataclass
from enum import StrEnum, auto
from pathlib import Path
import asdf
import numpy as np
from asdf.resource import DirectoryResourceMapping
from . import utils, version
# Register location of schema files
SCHEMA_PATH = Path(__file__).parent / "datamodel"
schema_mapping = DirectoryResourceMapping(
SCHEMA_PATH,
"asdf://cta-observatory.org/benchmark/",
recursive=True,
filename_pattern="*.yaml",
stem_filename=True,
)
asdf.get_config().add_resource_mapping(schema_mapping)
[docs]
class MetricStatus(StrEnum):
"""Reasons why a benchmark failed or passed."""
PASSED = auto()
FAILED_LOW = auto()
FAILED_HIGH = auto()
FAILED_OTHER = auto()
UNKNOWN = auto()
[docs]
@dataclass(slots=True)
class Storable:
"""Base class for metric data that can be inserted into a DataSet."""
name: str
schema_url = None
_defined_models = {} # Store all defined subclasses
def __init__(self):
# Default name
self.name = self.__class__.__name__
def __init_subclass__(cls, **kwargs):
"""Register every subclass for the 'open()' function."""
Storable._defined_models[cls.__name__] = cls
[docs]
@abstractmethod
def save(self, output_file: Path):
"""Serialize this object to output_file."""
[docs]
@classmethod
@abstractmethod
def load(cls, init):
"""Load this object from a file.
init : file path, file object, dict
- file path: Initialize from the given file (FITS, JSON or ASDF)
- readable file object: Initialize from the given file object
- dict: The object model tree for the data model
"""
def _init_treedict(self):
"""Prepare writing the file to asdf by creating the tree dictionary.
This function prepare the dictionary to ensure every
common parameter of all datamodels are set here (defined in core.schema).
Each sub-class need to class this, then add their own specificities to it.
Returns
-------
dict
Dictionary of all necessary information for the datamodel considered.
"""
if hasattr(datetime, "UTC"):
# New preferred way
now = datetime.datetime.now(datetime.UTC).isoformat()
else:
# deprecated in 3.12, but only solution in 3.10
now = datetime.datetime.utcnow().isoformat()
asdf_dict = {
"dm_module": self.__class__.__module__,
"filetype": self.__class__.__name__,
"name": self.name,
"date": now,
"testbench_version": version.__version__,
}
return asdf_dict
[docs]
def update_name_with_store(self, output_store):
"""Update the name of a storable.
Appends the name of a supplied output store.
"""
if output_store:
self.name = f"{output_store.name}"
[docs]
@dataclass
class ReqReport(Storable):
"""Stores the final report of a benchmark."""
schema_url = "asdf://cta-observatory.org/benchmark/report.schema"
value: float = np.nan
domain: str = "unknown" #: to which (telescope, array, etc) this report applies
status: MetricStatus = MetricStatus.UNKNOWN
[docs]
def save(self, output_file: Path):
"""Save object to file.
Parameters
----------
output_file : Path
Output filename.
"""
init = self._init_treedict()
init["value"] = self.value
init["domain"] = self.domain
init["status"] = self.status.value.upper()
asdffile = asdf.AsdfFile(init, custom_schema=self.schema_url)
asdffile.validate()
if not output_file.suffix == ".asdf":
output_file = output_file.with_suffix(output_file.suffix + ".asdf")
asdffile.write_to(output_file)
asdffile.close()
[docs]
@classmethod
def load(cls, init):
"""Load object from file or input dictionary.
Parameters
----------
init : dict
initialization parameters
"""
if isinstance(init, dict | asdf.AsdfFile):
asdffile = asdf.AsdfFile(tree=init, custom_schema=cls.schema_url)
elif isinstance(init, str | Path) or hasattr(init, "read"):
asdffile = asdf.open(init, custom_schema=cls.schema_url)
else:
raise ValueError(f"Unsupported type for `init`: '{init}'")
init_dict = utils.clean_tree(asdffile)
asdffile.close()
init_dict["status"] = getattr(MetricStatus, init_dict["status"])
obj = ReqReport(**init_dict)
return obj
def __repr_html__(self):
"""Html string representation."""
return f"<b>{self.title}</b>: {self.status}"
def open(init):
"""Create a DataModel from a number of different types.
Parameters
----------
init : file path, file object, dict
if file path: Initialize from the given file (FITS, JSON or ASDF),
if readable file object: Initialize from the given file object,
if dict: The object model tree for the data model.
Returns
-------
model : DataModel
data model instance
"""
if isinstance(init, Path):
init = str(init)
# If given a string, presume its a file path.
# if it has a read method, assume a file descriptor
if isinstance(init, str) or hasattr(init, "read"):
init = asdf.open(init)
# At that point, if init is not a dict, there's an issue
if not isinstance(init, dict | asdf.AsdfFile):
raise TypeError(f"Unsupported type for init: {type(init)}")
class_name = init["filetype"]
module_path = init["dm_module"]
# Make sure the class name is in _defined_models by forcing the dynamic use of that class
try:
getattr(importlib.import_module(module_path), class_name)
except ModuleNotFoundError:
msg = "Breaking change occurred in datapipe-testbench since you create this file. You need to create it again."
raise OSError(msg)
data = Storable._defined_models[class_name].load(dict(init))
init.close()
return data
