Update packages tests to write pickle and yaml

python/lsst/base/packages.py

@@ -287,21 +287,22 @@ def read(cls, filename):
         ----------
         filename : `str`
             Filename from which to read. The format is determined from the
-            file extension.  Currently support ``.pickle`` and ``.yaml``.
+            file extension.  Currently support ``.pickle``, ``.pkl``
+            and ``.yaml``.
 
         Returns
         -------
         packages : `Packages`
         """
         _, ext = os.path.splitext(filename)
-        if ext == ".pickle":
+        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)
         else:
-            raise ValueError(f"Unable to understand how to read file {filename}")
+            raise ValueError(f"Unable to determine how to read file {filename} from extension {ext}")
 
     def write(self, filename):
         """Write to file.
@@ -314,7 +315,7 @@ def write(self, filename):
             ``.pickle`` and ``.yaml``
         """
         _, ext = os.path.splitext(filename)
-        if ext == ".pickle":
+        if ext in (".pickle", ".pkl"):
             with open(filename, "wb") as ff:
                 pickle.dump(self, ff)
         elif ext == ".yaml":

tests/test_packages.py

@@ -62,20 +62,42 @@ def testRuntime(self):
         """
         lsst.base.getRuntimeVersions()
 
-    def testPackages(self):
-        """Test the Packages class"""
-        packages = lsst.base.Packages.fromSystem()
-
-        # Test pickling
+    def _writeTempFile(self, packages, suffix):
+        """Write packages to a temp file using the supplied suffix and read
+        back.
+        """
         # Can't use lsst.utils.tests.getTempFilePath because we're its dependency
-        temp = tempfile.NamedTemporaryFile(prefix="packages.", suffix=".pickle", delete=False)
+        temp = tempfile.NamedTemporaryFile(prefix="packages.", suffix=suffix, delete=False)
         tempName = temp.name
         temp.close()  # We don't use the fd, just want a filename
         try:
             packages.write(tempName)
             new = lsst.base.Packages.read(tempName)
         finally:
             os.unlink(tempName)
+        return new
+
+    def testPackages(self):
+        """Test the Packages class"""
+        packages = lsst.base.Packages.fromSystem()
+
+        # Test pickling and YAML
+        new = self._writeTempFile(packages, ".pickle")
+        new_pkl = self._writeTempFile(packages, ".pkl")
+        new_yaml = self._writeTempFile(packages, ".yaml")
+
+        self.assertIsInstance(new, lsst.base.Packages)
+        self.assertIsInstance(new_yaml, lsst.base.Packages)
+        self.assertEqual(new, packages)
+        self.assertEqual(new_pkl, new)
+        self.assertEqual(new, new_yaml)
+
+        with self.assertRaises(ValueError):
+            self._writeTempFile(packages, ".txt")
+
+        with self.assertRaises(ValueError):
+            # .txt extension is not recognized
+            lsst.base.Packages.read("something.txt")
 
         # 'packages' and 'new' should have identical content
         self.assertDictEqual(packages.difference(new), {})