Merge pull request #4 from bauman/nopy2

fixup py3, remove py2

Author: bauman

bson2json.py

@@ -4,8 +4,6 @@
 from bsonstream import BSONInput
 from sys import argv, stdout
 import gzip
-import six
-
 
 if "gz" in argv[1] or "dz" in argv[1]:
     f = gzip.open(argv[1], 'rb')
@@ -15,8 +13,5 @@
 stream = BSONInput(fh=f)
 for doc in stream:
     json_str = dumps(doc)
-    if six.PY3:
-        stdout.write("{}\n".format(json_str))
-    else:  # six.PY2
-        stdout.write("%s\n" %json_str)
+    stdout.write(f"{json_str}\n")
 

bsonstream/__init__.py

@@ -1,5 +1,4 @@
 from bson import InvalidBSON, BSON, codec_options
-import six
 import sys
 import struct
 
@@ -15,9 +14,14 @@ class BSONInput(object):
     https://github.com/klbostee/typedbytes
     """
 
-    def __init__(self, fh=sys.stdin, unicode_errors='strict', fast_string_prematch="", decode=True):
+    def __init__(self, fh=sys.stdin,
+                 unicode_errors: str = 'strict',
+                 fast_string_prematch: bytes = b"",
+                 decode: bool = True):
         self.fh = fh
         self.unicode_errors = unicode_errors
+        if not isinstance(fast_string_prematch, bytes):
+            raise ValueError("fast_string_prematch must be a bytes object")
         self.fast_string_prematch=fast_string_prematch
         self.eof = False
         self.decode = decode
@@ -26,26 +30,22 @@ def __init__(self, fh=sys.stdin, unicode_errors='strict', fast_string_prematch="
     def _read(self):
         try:
             size_bits = self.fh.read(4)
-            size = struct.unpack("<i", size_bits)[0] - 4 # BSON size byte includes itself 
+            size = struct.unpack("<i", size_bits)[0] - 4  # BSON size byte includes itself
             data = size_bits + self.fh.read(size)
             if len(data) != size + 4:
                 raise struct.error("Unable to read expected BSON Chunk; " +
                                    "EOF, underful buffer or invalid object size.")
-            if six.PY3:
-                eoo = 0x00
-            else:  # six.PY2
-                eoo = "\x00"
+            eoo = 0x00
             if data[size + 4 - 1] != eoo:
                 raise InvalidBSON("Bad EOO in BSON Data")
-            if self.fast_string_prematch.encode("utf-8") in data:
+            if self.fast_string_prematch in data:
                 if self.decode:
                     try:
                         return BSON(data).decode(self.codec)
                     except TypeError:
                         return BSON(data).decode()
                 else:
                     return data
-            raise ValueError("Unknown Error")
         except struct.error as e:
             self.eof = True
             raise StopIteration(e)
@@ -60,13 +60,19 @@ def read(self):
     def _reads(self):
         r = self._read
         while 1:
-            yield r()
+            try:
+                data = r()
+                if data:
+                    yield data
+            except StopIteration:
+                break
 
     def close(self):
         self.fh.close()
 
     __iter__ = reads = _reads
 
+
 class KeyValueBSONInput(BSONInput):
     def read(self):
         try:
@@ -78,12 +84,12 @@ def read(self):
 
     def reads(self):
         it = self._reads()
-        if six.PY3:
-            n = it.__next__
-        else:  # six.PY2
-            n = it.next
+        n = it.__next__
         while 1:
-            doc = n()
-            yield doc
+            try:
+                doc = n()
+                yield doc
+            except StopIteration:
+                break
 
     __iter__ = reads