Skip to content

Commit

Permalink
Rewrite serialization to allow access to byte form
Browse files Browse the repository at this point in the history
This lets you bypass a file completely if using an
object store.
  • Loading branch information
timj committed Jun 11, 2020
1 parent 04f1aa2 commit 17ab1bc
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 15 deletions.
66 changes: 51 additions & 15 deletions python/lsst/base/packages.py
Original file line number Diff line number Diff line change
Expand Up @@ -257,6 +257,10 @@ class Packages:
This is essentially a wrapper around a dict with some conveniences.
"""

formats = {".pkl": "pickle",
".pickle": "pickle",
".yaml": "yaml"}

def __init__(self, packages):
assert isinstance(packages, Mapping)
self._packages = packages
Expand All @@ -279,6 +283,24 @@ def fromSystem(cls):
packages.update(getEnvironmentPackages()) # Should be last, to override products with LOCAL versions
return cls(packages)

@classmethod
def fromBytes(cls, data, format):
"""Construct the object from a byte representation.
Parameters
----------
data : `bytes`
The serialized form of this object in bytes.
format : `str`
The format of those bytes. Can be ``yaml`` or ``pickle``.
"""
if format == "pickle":
return pickle.loads(data)
elif format == "yaml":
return yaml.load(data, Loader=yaml.SafeLoader)
else:
raise ValueError(f"Unexpected serialization format given: {format}")

@classmethod
def read(cls, filename):
"""Read packages from filename.
Expand All @@ -295,14 +317,32 @@ def read(cls, filename):
packages : `Packages`
"""
_, ext = os.path.splitext(filename)
if ext in (".pickle", ".pkl"):
with open(filename, "rb") as ff:
return pickle.load(ff)
elif ext == ".yaml":
with open(filename, "r") as ff:
return yaml.load(ff, Loader=yaml.SafeLoader)
if ext not in cls.formats:
raise ValueError(f"Format from {ext} extension in file {filename} not recognized")
with open(filename, "rb") as ff:
data = ff.read()
return cls.fromBytes(data, cls.formats[ext])

def toBytes(self, format):
"""Convert the object to a serialized bytes form using the
specified format.
Parameters
----------
format : `str`
Format to use when serializing. Can be ``yaml`` or ``pickle``.
Returns
-------
data : `bytes`
Byte string representing the serialized object.
"""
if format == "pickle":
return pickle.dumps(self)
elif format == "yaml":
return yaml.dump(self).encode("utf-8")
else:
raise ValueError(f"Unable to determine how to read file {filename} from extension {ext}")
raise ValueError(f"Unexpected serialization format requested: {format}")

def write(self, filename):
"""Write to file.
Expand All @@ -315,14 +355,10 @@ def write(self, filename):
``.pickle`` and ``.yaml``
"""
_, ext = os.path.splitext(filename)
if ext in (".pickle", ".pkl"):
with open(filename, "wb") as ff:
pickle.dump(self, ff)
elif ext == ".yaml":
with open(filename, "w") as ff:
yaml.dump(self, ff)
else:
raise ValueError(f"Unexpected file format requested: {ext}")
if ext not in self.formats:
raise ValueError(f"Format from {ext} extension in file {filename} not recognized")
with open(filename, "wb") as ff:
ff.write(self.toBytes(self.formats[ext]))

def __len__(self):
return len(self._packages)
Expand Down
12 changes: 12 additions & 0 deletions tests/test_packages.py
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,18 @@ def testPackages(self):
self.assertDictEqual(packages.extra(new), {})
self.assertEqual(len(packages), len(new))

# Serialize via bytes
for format in ("pickle", "yaml"):
asbytes = new.toBytes(format)
from_bytes = lsst.base.Packages.fromBytes(asbytes, format)
self.assertEqual(from_bytes, new)

with self.assertRaises(ValueError):
new.toBytes("unknown_format")

with self.assertRaises(ValueError):
lsst.base.Packages.fromBytes(from_bytes, "unknown_format")


if __name__ == "__main__":
unittest.main()

0 comments on commit 17ab1bc

Please sign in to comment.