From 79151ac7a489a4684d5b9fe3a2e581e64ae8c24a Mon Sep 17 00:00:00 2001
From: James Davies <jdavies@stsci.edu>
Date: Fri, 7 Dec 2018 11:25:24 -0500
Subject: [PATCH] Make SPK class close mmaps when exiting context manager (#25)

* Properly close mmaps; respect with context manager
* Update test to use with context manager
* Remove try/except in SPK.close()
* Restore SPK mmap attributes to None on close()
---
 jplephem/spk.py  | 10 +++++++++-
 jplephem/test.py | 13 ++++++-------
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/jplephem/spk.py b/jplephem/spk.py
index 3067604..e8819ac 100644
--- a/jplephem/spk.py
+++ b/jplephem/spk.py
@@ -48,7 +48,9 @@ class SPK(object):
         self.daf.file.close()
         for segment in self.segments:
             if hasattr(segment, '_data'):
-                del segment._data  # TODO: explicitly close each memory map
+                del segment._data
+        self.daf._array = None
+        self.daf._map = None
 
     def __str__(self):
         daf = self.daf
@@ -65,6 +67,12 @@ class SPK(object):
         """Return the file comments, as a string."""
         return self.daf.comments()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
 
 class Segment(object):
     """A single segment of an SPK file.
diff --git a/jplephem/test.py b/jplephem/test.py
index d4e5243..b9827f5 100644
--- a/jplephem/test.py
+++ b/jplephem/test.py
@@ -365,13 +365,12 @@ class LegacyTests(_CommonTests, TestCase):
 class NAIF_DAF_Tests(TestCase):
 
     def test_single_position(self):
-        kernel = SPK(NAIF_DAF(open('de405.bsp', 'rb')))
-        x, y, z = kernel[0,4].compute(2457061.5)
-        # Expect rough agreement with a DE430 position from our README:
-        self.assertAlmostEqual(x, 2.05700211e+08, delta=2.0)
-        self.assertAlmostEqual(y, 4.25141646e+07, delta=2.0)
-        self.assertAlmostEqual(z, 1.39379183e+07, delta=2.0)
-        kernel.close()
+        with SPK(NAIF_DAF(open('de405.bsp', 'rb'))) as kernel:
+            x, y, z = kernel[0,4].compute(2457061.5)
+            # Expect rough agreement with a DE430 position from our README:
+            self.assertAlmostEqual(x, 2.05700211e+08, delta=2.0)
+            self.assertAlmostEqual(y, 4.25141646e+07, delta=2.0)
+            self.assertAlmostEqual(z, 1.39379183e+07, delta=2.0)
 
 
 class CommandLineTests(TestCase):