# -*- coding: UTF-8 -*-
"""Classes and functions to support working with alerts and related data.
.. autosummary::
Cast
----
"""
import base64
import collections
import io
import json
import logging
import astropy.table
import astropy.time
import attrs
import fastavro
LOGGER = logging.getLogger(__name__)
[docs]
@attrs.define
class Cast:
"""Methods to convert data types."""
[docs]
@staticmethod
def bytes_to_b64utf8(bytes_data):
"""Convert bytes data to UTF-8.
Args:
bytes_data (bytes):
Data to be converted to UTF-8.
Returns:
str:
The ``bytes_data`` converted to a string in UTF-8 format.
"""
if bytes_data is not None:
return base64.b64encode(bytes_data).decode("utf-8")
[docs]
@staticmethod
def json_to_dict(bytes_data):
"""Converts JSON serialized bytes data to a dictionary.
Args:
bytes_data (bytes):
Data to be converted to a dictionary.
Returns:
dict:
The unpacked dictionary from the ``bytes_data``.
"""
if bytes_data is not None:
return json.loads(bytes_data)
[docs]
@staticmethod
def b64json_to_dict(bytes_data):
"""Converts base64 encoded, JSON serialized bytes data to a dictionary.
Args:
bytes_data (Base64):
Data to be converted to a dictionary.
Returns:
dict:
The unpacked dictionary from the ``bytes_data``.
"""
if bytes_data is not None:
return Cast.json_to_dict(base64.b64decode(bytes_data))
[docs]
@staticmethod
def avro_to_dict(bytes_data):
"""Converts Avro serialized bytes data to a dictionary.
Args:
bytes_data (bytes):
Avro serialized bytes data to be converted to a dictionary. The schema must be attached in the header.
Returns:
dict:
The unpacked dictionary from the ``bytes_data``.
"""
if bytes_data is not None:
with io.BytesIO(bytes_data) as fin:
alert_dicts = list(fastavro.reader(fin)) # list with single dict
if len(alert_dicts) != 1:
LOGGER.warning(f"Expected 1 Avro record. Found {len(alert_dicts)}.")
return alert_dicts[0]
[docs]
@staticmethod
def b64avro_to_dict(bytes_data):
"""Converts base64 encoded, Avro serialized bytes data to a dictionary.
Args:
bytes_data (bytes):
Base64 encoded, Avro serialized bytes data to be converted to a dictionary.
Returns:
dict:
The unpacked dictionary from the ``bytes_data``.
"""
return Cast.avro_to_dict(base64.b64decode(bytes_data))
# --- Work with alert dictionaries
[docs]
@staticmethod
def alert_dict_to_table(alert_dict: dict) -> astropy.table.Table:
"""Package a ZTF alert dictionary into an Astropy Table.
Args:
alert_dict (dict):
A dictionary containing ZTF alert information.
Returns:
astropy.table.Table:
An Astropy Table containing the alert information.
"""
# collect rows for the table
candidate = collections.OrderedDict(alert_dict["candidate"])
rows = [candidate]
for prv_cand in alert_dict["prv_candidates"]:
# astropy 3.2.1 cannot handle dicts with different keys (fixed by 4.1)
prv_cand_tmp = {key: prv_cand.get(key, None) for key in candidate.keys()}
rows.append(prv_cand_tmp)
# create and return the table
table = astropy.table.Table(rows=rows)
table.meta["comments"] = f"ZTF objectId: {alert_dict['objectId']}"
return table
@staticmethod
def _strip_cutouts_ztf(alert_dict: dict) -> dict:
"""Drop the cutouts from the alert dictionary.
Args:
alert_dict (dict):
ZTF alert formatted as a dictionary.
Returns:
dict:
The modified `alert_dict` with the cutouts (postage stamps) removed.
"""
cutouts = ["cutoutScience", "cutoutTemplate", "cutoutDifference"]
alert_stripped = {k: v for k, v in alert_dict.items() if k not in cutouts}
return alert_stripped
# dates
[docs]
@staticmethod
def jd_to_readable_date(jd) -> str:
"""Converts a Julian date to a human-readable string.
Args:
jd (float):
Datetime value in Julian format.
Returns:
str:
The ``jd`` in the format 'day mon year hour:min'.
"""
return astropy.time.Time(jd, format="jd").strftime("%d %b %Y - %H:%M:%S")
# --- Survey-specific
[docs]
def ztf_fid_names() -> dict:
"""Return a dictionary mapping the ZTF `fid` (filter ID) to the common name."""
return {1: "g", 2: "r", 3: "i"}