Coverage for python/lsst/meas/transiNet/modelPackages/formatters.py: 42%
22 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-23 03:57 -0700
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-23 03:57 -0700
1from lsst.daf.butler import Formatter
2from io import BytesIO
4__all__ = ["NNModelPackageFormatter", "NNModelPackagePayload"]
7class NNModelPackagePayload():
8 """A thin wrapper around the payload of a NNModelPackageFormatter,
9 which simply carries an in-memory file between the formatter and the
10 storage adapter of model pacakges.
11 """
12 def __init__(self):
13 self.bytes = BytesIO()
16class NNModelPackageFormatter(Formatter):
17 """Formatter for NN model packages.
18 """
19 extension = ".zip"
21 def read(self, component=None):
22 """Read a dataset.
24 Parameters
25 ----------
26 component : `str`, optional
27 Component to read from the file.
29 Returns
30 -------
31 payload : `NNModelPackagePayload`
32 The requested data as a Python object.
33 """
34 payload = NNModelPackagePayload()
35 with open(self.fileDescriptor.location.path, "rb") as f:
36 payload.bytes = BytesIO(f.read())
37 return payload
39 def write(self, inMemoryDataset):
40 """Write a Dataset.
42 Parameters
43 ----------
44 inMemoryDataset : `object`
45 The Dataset to store.
47 """
48 with open(self.fileDescriptor.location.path, "wb") as f:
49 f.write(inMemoryDataset.bytes.getvalue())
51 def fromBytes(self, serializedDataset: bytes, component=None):
52 """Read serialized data into a Dataset or its component.
54 Parameters
55 ----------
56 serializedDataset : `bytes`
57 Bytes object to unserialize.
58 component : `str`, optional
59 Component to read from the Dataset. Only used if the `StorageClass`
60 for reading differed from the `StorageClass` used to write the
61 file.
63 Returns
64 -------
65 payload : `NNModelPackagePayload`
66 The requested data as a Python object.
67 """
68 payload = NNModelPackagePayload()
69 payload.bytes = BytesIO(serializedDataset)
70 return payload
72 def toBytes(self, inMemoryDataset):
73 """Serialize the Dataset to bytes based on formatter.
75 Parameters
76 ----------
77 inMemoryDataset : `object`
78 The Python object to serialize.
80 Returns
81 -------
82 serializedDataset : `bytes`
83 Bytes representing the serialized dataset.
84 """
85 return inMemoryDataset.bytes.getvalue()