from io import BytesIO
from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Optional
from uuid import UUID
from google.protobuf.internal.containers import MessageMap
from .node import Node
from .proto import AuxData_pb2
from .serialization import Serialization
from .util import DictLike
if TYPE_CHECKING: # pragma: no cover
# Ignore flake8 "imported but unused" errors.
from .ir import IR # noqa: F401
class _LazyDataContainer:
"""
Container that holds the raw byte stream until it is read, then releases
it. If it is never read, then serialization skips re-encoding (and
deserializing) the data.
"""
def __init__(
self,
raw_data: bytes,
type_name: str,
get_by_uuid: Callable[[UUID], Optional[Node]],
):
self.raw_data: Optional[bytes] = raw_data
self.type_name = type_name
self.get_by_uuid = get_by_uuid
def get_data(self) -> object:
"""
Get any pending still-serialized data, or return the passed data
instead (the default).
"""
assert self.raw_data is not None
rv = AuxData.serializer.decode(
self.raw_data, self.type_name, self.get_by_uuid
)
self.raw_data = None
return rv
def get_raw_data(self) -> bytes:
""" """
assert self.raw_data is not None
return self.raw_data
[docs]class AuxData:
"""AuxData objects can be attached to the :class:`gtirb.IR` or individual
:class:`gtirb.Module` s to store additional client-specific data in a
portable way.
AuxData represents a portable, language-independent manner of encoding
rich data. To do this, all data is stored on disk as a series of bytes
with a string describing the format of the data, called a *type name*. See
:mod:`gtirb.serialization` for the list of all default types. Types may
also be parameterized; for example, ``mapping<string,UUID>`` is a ``dict``
from ``str`` objects to ``UUID`` objects. All ``AuxData`` requires
a valid type name in order to be serialized.
:ivar ~.data: The value stored in this AuxData.
:ivar ~.type_name: A string describing the type of ``data``.
Used to determine the proper codec for serializing this AuxData.
"""
serializer: ClassVar[Serialization] = Serialization()
"""This is a :class:`gtirb.Serialization` instance, used to
encode and decode ``data`` fields of all ``AuxData``. See
:mod:`gtirb.serialization` for details.
"""
[docs] def __init__(
self,
data: object,
type_name: str,
lazy_container: Optional[_LazyDataContainer] = None,
):
"""
:param data: The value stored in this AuxData.
:param type_name: A string describing the type of ``data``.
Used to determine the proper codec for serializing this AuxData.
:param lazy_container: An object that will lazily deserialize the
auxdata table backing this object, or None.
"""
self._lazy_container = lazy_container
# _data has type Any to avoid disrupting clients want to type check
# their use of gtirb. If _data had type object, they would have to
# verify the element types of potentially large containers, or else
# just subvert the type system by casting anyway.
self._data: Any = data # type: ignore[misc]
self.type_name = type_name
@property
def data(self) -> Any: # type: ignore[misc]
if self._lazy_container is not None:
self._data = self._lazy_container.get_data()
self._lazy_container = None
return self._data
@data.setter
def data(self, value: object) -> None:
self._data = value
self._lazy_container = None
@classmethod
def _from_protobuf(
cls,
aux_data: AuxData_pb2.AuxData,
ir: Optional["IR"],
) -> "AuxData":
"""Deserialize AuxData from Protobuf. Lazy, will not perform
deserialization until .data is accessed.
:param aux_data: The Protobuf AuxData object.
"""
# Defer deserialization until someone accesses .data
assert ir
lazy_container = _LazyDataContainer(
aux_data.data, aux_data.type_name, ir.get_by_uuid
)
return cls(
data=None,
type_name=aux_data.type_name,
lazy_container=lazy_container,
)
def _to_protobuf(self) -> AuxData_pb2.AuxData:
"""Get a Protobuf representation of the AuxData."""
proto_auxdata = AuxData_pb2.AuxData()
proto_auxdata.type_name = self.type_name
# If we are serializing the same data, and the way that data is encoded
# has not changed, then just use the already serialized copy.
if self._lazy_container is not None and (
self.type_name == self._lazy_container.type_name
):
proto_auxdata.data = self._lazy_container.get_raw_data()
else:
data_stream = BytesIO()
AuxData.serializer.encode(data_stream, self.data, self.type_name)
proto_auxdata.data = data_stream.getvalue()
return proto_auxdata
def __repr__(self) -> str:
return (
"AuxData("
"type_name={type_name!r}, "
"data={data!r}, "
")".format(type_name=self.type_name, data=self.data)
)
[docs]class AuxDataContainer(Node):
"""The base class for anything that holds AuxData tables; that is,
:class:`gtirb.IR` and :class:`gtirb.Module`.
:ivar ~.aux_data: The auxiliary data associated
with the object, as a mapping from names to
:class:`gtirb.AuxData`.
"""
[docs] def __init__(
self,
aux_data: DictLike[str, AuxData] = {},
uuid: Optional[UUID] = None,
):
"""
:param aux_data: The initial auxiliary data to be associated
with the object, as a mapping from names to
:class:`gtirb.AuxData`. Defaults to an empty :class:`dict`.
:param uuid: the UUID of this ``AuxDataContainer``,
or None if a new UUID needs generated via :func:`uuid.uuid4`.
Defaults to None.
"""
super().__init__(uuid)
self.aux_data: Dict[str, AuxData] = dict(aux_data)
@classmethod
def _read_protobuf_aux_data(
cls,
proto_container: "MessageMap[str, AuxData_pb2.AuxData]",
ir: Optional["IR"],
) -> Dict[str, AuxData]:
"""
Instead of the overrided _decode_protobuf, this method requires the
Protobuf message to read from. AuxDataContainers need to call this
method in their own _decode_protobuf overrides.
:param proto_container: A Protobuf message with a field called
``aux_data``.
"""
return {
key: AuxData._from_protobuf(val, ir)
for key, val in proto_container.items()
}
def _write_protobuf_aux_data(
self, proto_container: "MessageMap[str, AuxData_pb2.AuxData]"
) -> None:
"""
Instead of the overrided _to_protobuf, this method requires the
Protobuf message to write into. AuxDataContainers need to call this
method in their own _to_protobuf overrides.
:param proto_container: A Protobuf message with a field called
``aux_data``.
"""
for k, v in self.aux_data.items():
proto_container[k].CopyFrom(v._to_protobuf())
[docs] def deep_eq(self, other: object) -> bool:
"""This overrides :func:`gtirb.Node.deep_eq` to check for
AuxData equality.
Because the values stored by AuxData are not necessarily
amenable to deep checking, the auxiliary data dictionaries
stored for ``self`` and ``other`` are not deeply checked. Instead,
they are considered to be equal if their sets of keys are equal.
"""
if not isinstance(other, AuxDataContainer):
return False
if (
self.uuid != other.uuid
or self.aux_data.keys() != other.aux_data.keys()
):
return False
return True