Coverage for python / lsst / meas / transiNet / modelPackages / formatters.py: 71%

1from lsst.daf.butler import FormatterV2 

2from lsst.resources import ResourcePath 

3from io import BytesIO 

4from typing import Any 

5 

6__all__ = ["NNModelPackageFormatter", "NNModelPackagePayload"] 

7 

8 

9class NNModelPackagePayload(): 

10 """A thin wrapper around the payload of a NNModelPackageFormatter, 

11 which simply carries an in-memory file between the formatter and the 

12 storage adapter of model pacakges. 

13 """ 

14 def __init__(self): 

15 self.bytes = BytesIO() 

16 

17 

18class NNModelPackageFormatter(FormatterV2): 

19 """Formatter for NN model packages. 

20 """ 

21 default_extension = ".zip" 

22 can_read_from_uri = True 

23 

24 def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any: 

25 """Read a dataset. 

26 

27 Parameters 

28 ---------- 

29 uri : `lsst.ResourcePath` 

30 Location of the file to read. 

31 component : `str` or `None`, optional 

32 Component to read from the file. 

33 expected_size : `int`, optional 

34 Expected size of the file. 

35 

36 Returns 

37 ------- 

38 payload : `NNModelPackagePayload` 

39 The requested data as a Python object. 

40 """ 

41 payload = NNModelPackagePayload() 

42 payload.bytes = BytesIO(uri.read()) 

43 return payload 

44 

45 def to_bytes(self, in_memory_dataset: Any) -> bytes: 

46 return in_memory_dataset.bytes.getvalue()