clean up nmea_file a bit

Knio · Sep 29, 2015 · d98c056 · d98c056
1 parent 2af0543
commit d98c056
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 35 deletions.
diff --git a/pynmea2/nmea_file.py b/pynmea2/nmea_file.py
@@ -1,15 +1,24 @@
-from __future__ import unicode_literals
-from pynmea2.nmea import NMEASentence
+try:
+    # pylint: disable=used-before-assignment
+    basestring = basestring
+except NameError: # py3
+    basestring = str
+
+from .nmea import NMEASentence
 
 
 class NMEAFile(object):
     """
     Reads NMEA sentences from a file similar to a standard python file object.
     """
 
-    def __init__(self, path, mode='r'):
+    def __init__(self, f, *args, **kwargs):
         super(NMEAFile, self).__init__()
-        self._file = self.open(path, mode=mode)
+        if isinstance(f, basestring) or args or kwargs:
+            self._file = self.open(f, *args, **kwargs)
+        else:
+            self._file = f
+        self._context = None
 
     def open(self, fp, mode='r'):
         """
@@ -33,11 +42,15 @@ def __iter__(self):
             yield self.parse(line)
 
     def __enter__(self):
-        self._file.__enter__()
+        if hasattr(self._file, '__enter__'):
+            self._context = self._file.__enter__()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
+        if self._context:
+            ctx = self._context
+            self._context = None
+            ctx.__exit__(exc_type, exc_val, exc_tb)
 
     def next(self):
         """

diff --git a/pynmea2/types/proprietary/ubx.py b/pynmea2/types/proprietary/ubx.py
@@ -1,11 +1,11 @@
-# u-blox
-
 # pylint: disable=wildcard-import,unused-wildcard-import
+from decimal import Decimal
+
 from ... import nmea
 from ...nmea_utils import *
-from decimal import Decimal
 
 
+# u-blox
 class UBX(nmea.ProprietarySentence):
     sentence_types = {}
 

diff --git a/test/test_file.py b/test/test_file.py
@@ -4,52 +4,42 @@
     from io import StringIO
 
 import pynmea2
-import os
-
-
 
+TEST_DATA = """$GPRMC,181031.576,V,3926.276,N,07739.361,W,99.7,18.30,250915,,E*79
+$GPGGA,181032.576,3926.276,N,07739.361,W,0,00,,,M,,M,,*5F
+$GPGLL,3926.276,N,07739.361,W,181033.576,V*3E
+$GPRMC,181034.576,V,3949.797,N,07809.854,W,18.8,34.66,250915,,E*75
+$GPGGA,181035.576,3949.797,N,07809.854,W,0,00,,,M,,M,,*5A
+$GPGLL,3949.797,N,07809.854,W,181036.576,V*39
+$GPRMC,181037.576,V,4040.018,N,07808.022,W,32.9,16.43,250915,,E*77
+$GPGGA,181038.576,4040.018,N,07808.022,W,0,00,,,M,,M,,*58
+$GPGLL,4040.018,N,07808.022,W,181039.576,V*39
+$GPRMC,181040.576,V,4133.618,N,07725.034,W,96.8,44.47,250915,,E*7F"""
 
 def test_file():
-
-    filepath = './nmea_file'
-
-    with open(filepath, 'w') as f:
-        f.write("""$GPRMC,181031.576,V,3926.276,N,07739.361,W,99.7,18.30,250915,,E*79
-    $GPGGA,181032.576,3926.276,N,07739.361,W,0,00,,,M,,M,,*5F
-    $GPGLL,3926.276,N,07739.361,W,181033.576,V*3E
-    $GPRMC,181034.576,V,3949.797,N,07809.854,W,18.8,34.66,250915,,E*75
-    $GPGGA,181035.576,3949.797,N,07809.854,W,0,00,,,M,,M,,*5A
-    $GPGLL,3949.797,N,07809.854,W,181036.576,V*39
-    $GPRMC,181037.576,V,4040.018,N,07808.022,W,32.9,16.43,250915,,E*77
-    $GPGGA,181038.576,4040.018,N,07808.022,W,0,00,,,M,,M,,*58
-    $GPGLL,4040.018,N,07808.022,W,181039.576,V*39
-    $GPRMC,181040.576,V,4133.618,N,07725.034,W,96.8,44.47,250915,,E*7F"""
-                    )
-
-    nmeafile = pynmea2.NMEAFile(filepath)
+    nmeafile = pynmea2.NMEAFile(StringIO(TEST_DATA))
 
     nmea_strings = nmeafile.read()
     assert len(nmea_strings) == 10
     assert all([isinstance(s, pynmea2.NMEASentence) for s in nmea_strings])
     del nmeafile
 
-    with pynmea2.NMEAFile(filepath) as _f:
+    with pynmea2.NMEAFile(StringIO(TEST_DATA)) as _f:
         nmea_strings = [_f.readline() for i in range(10)]
     assert len(nmea_strings) == 10
     assert all([isinstance(s, pynmea2.NMEASentence) for s in nmea_strings])
 
-    with pynmea2.NMEAFile(filepath) as _f:
+    with pynmea2.NMEAFile(StringIO(TEST_DATA)) as _f:
         nmea_strings = [s for s in _f]
     assert len(nmea_strings) == 10
     assert all([isinstance(s, pynmea2.NMEASentence) for s in nmea_strings])
 
-    with pynmea2.NMEAFile(filepath) as _f:
+    with pynmea2.NMEAFile(StringIO(TEST_DATA)) as _f:
         nmea_strings = [_f.next() for i in range(10)]
     assert len(nmea_strings) == 10
     assert all([isinstance(s, pynmea2.NMEASentence) for s in nmea_strings])
 
-    if os.path.isfile(filepath):
-        os.remove(filepath)
 
-test_file()
+if __name__ == '__main__':
+    test_file()