changed interface for getting default tokenizer

Author: daggaz

src/json_stream/httpx/__init__.py

@@ -1,5 +1,4 @@
 import json_stream
-from json_stream.select_tokenizer import default_tokenizer
 
 
 CONTENT_CHUNK_SIZE = 10 * 1024
@@ -9,9 +8,21 @@ def _to_iterable(response, chunk_size):
     return response.iter_bytes(chunk_size=chunk_size)
 
 
-def load(response, persistent=False, tokenizer=default_tokenizer, chunk_size=CONTENT_CHUNK_SIZE):
-    return json_stream.load(_to_iterable(response, chunk_size), persistent=persistent, tokenizer=tokenizer)
-
-
-def visit(response, visitor, tokenizer=default_tokenizer, chunk_size=CONTENT_CHUNK_SIZE):
-    return json_stream.visit(_to_iterable(response, chunk_size), visitor, tokenizer=tokenizer)
+def load(response, persistent=False, tokenizer=None, chunk_size=CONTENT_CHUNK_SIZE, buffering=0, **kwargs):
+    return json_stream.load(
+        _to_iterable(response, chunk_size),
+        persistent=persistent,
+        tokenizer=tokenizer,
+        buffering=buffering,
+        **kwargs
+    )
+
+
+def visit(response, visitor, tokenizer=None, chunk_size=CONTENT_CHUNK_SIZE, buffering=0, **kwargs):
+    return json_stream.visit(
+        _to_iterable(response, chunk_size),
+        visitor,
+        tokenizer=tokenizer,
+        buffering=buffering,
+        **kwargs
+    )

src/json_stream/loader.py

@@ -1,11 +1,9 @@
 from json_stream.base import StreamingJSONBase, TokenType
-from json_stream.iterators import ensure_file
-from json_stream.select_tokenizer import default_tokenizer
+from json_stream.select_tokenizer import get_token_stream
 
 
-def load(fp_or_iterable, persistent=False, tokenizer=default_tokenizer, buffering=-1):
-    fp = ensure_file(fp_or_iterable)
-    token_stream = tokenizer(fp, buffering=buffering)
+def load(fp_or_iterable, persistent=False, tokenizer=None, buffering=-1, **kwargs):
+    token_stream = get_token_stream(fp_or_iterable, tokenizer=tokenizer, buffering=buffering, **kwargs)
     token_type, token = next(token_stream)
     if token_type == TokenType.OPERATOR:
         return StreamingJSONBase.factory(token, token_stream, persistent)

src/json_stream/requests/__init__.py

@@ -1,5 +1,4 @@
 import json_stream
-from json_stream.select_tokenizer import default_tokenizer
 
 
 CONTENT_CHUNK_SIZE = 10 * 1024
@@ -9,9 +8,23 @@ def _to_iterable(response, chunk_size):
     return response.iter_content(chunk_size=chunk_size)
 
 
-def load(response, persistent=False, tokenizer=default_tokenizer, chunk_size=CONTENT_CHUNK_SIZE):
-    return json_stream.load(_to_iterable(response, chunk_size), persistent=persistent, tokenizer=tokenizer)
-
-
-def visit(response, visitor, tokenizer=default_tokenizer, chunk_size=CONTENT_CHUNK_SIZE):
-    return json_stream.visit(_to_iterable(response, chunk_size), visitor, tokenizer=tokenizer)
+def load(response, persistent=False, tokenizer=None, chunk_size=CONTENT_CHUNK_SIZE,
+         buffering=0, **kwargs):
+    return json_stream.load(
+        _to_iterable(response, chunk_size),
+        persistent=persistent,
+        tokenizer=tokenizer,
+        buffering=buffering,
+        **kwargs
+    )
+
+
+def visit(response, visitor, tokenizer=None, chunk_size=CONTENT_CHUNK_SIZE,
+          buffering=0, **kwargs):
+    return json_stream.visit(
+        _to_iterable(response, chunk_size),
+        visitor,
+        tokenizer=tokenizer,
+        buffering=buffering,
+        **kwargs
+    )

src/json_stream/select_tokenizer.py

@@ -1,12 +1,20 @@
 from warnings import warn
 
+from json_stream.iterators import ensure_file
 from json_stream.tokenizer import tokenize
 from json_stream_rs_tokenizer import rust_tokenizer_or_raise, ExtensionException
 
-try:
-    default_tokenizer = rust_tokenizer_or_raise()
-except ExtensionException as e:
-    warn(str(e), category=ImportWarning)  # ImportWarnings are ignored by default
-    default_tokenizer = tokenize
 
-__all__ = ['default_tokenizer']
+def get_tokenizer(**kwargs):
+    try:
+        return rust_tokenizer_or_raise(**kwargs)
+    except ExtensionException as e:
+        warn(str(e), category=ImportWarning)  # ImportWarnings are ignored by default
+        return tokenize
+
+
+def get_token_stream(fp_or_iterable, tokenizer, **tokenizer_kwargs):
+    fp = ensure_file(fp_or_iterable)
+    if tokenizer is None:
+        tokenizer = get_tokenizer(**tokenizer_kwargs)
+    return tokenizer(fp, **tokenizer_kwargs)

src/json_stream/tests/__init__.py

@@ -3,14 +3,13 @@
 from itertools import zip_longest
 from unittest import TestCase
 
-from json_stream.select_tokenizer import default_tokenizer
 
 from json_stream import load
 from json_stream.base import TransientAccessException
 
 
 class JSONLoadTestCase(TestCase):
-    def _test_object(self, obj, persistent, binary=False, tokenizer=default_tokenizer):
+    def _test_object(self, obj, persistent, binary=False, tokenizer=None):
         self.assertListEqual(list(self._to_data(obj, persistent, binary, tokenizer)), list(obj))
         self.assertListEqual(list(self._to_data(obj, persistent, binary, tokenizer).keys()), list(obj.keys()))
         self.assertListEqual(list(self._to_data(obj, persistent, binary, tokenizer).values()), list(obj.values()))
@@ -40,7 +39,7 @@ def _test_object(self, obj, persistent, binary=False, tokenizer=default_tokenize
             with self.assertRaises(TransientAccessException):
                 data.items()  # can't get keys
 
-    def _test_list(self, obj, persistent, binary=False, tokenizer=default_tokenizer):
+    def _test_list(self, obj, persistent, binary=False, tokenizer=None):
         self.assertListEqual(list(self._to_data(obj, persistent, binary, tokenizer)), list(obj))
         if persistent:
             self.assertEqual(len(self._to_data(obj, persistent, binary, tokenizer)), len(obj))

src/json_stream/tests/test_tokenizer_integration.py

@@ -5,20 +5,20 @@
 
 from json_stream import load
 
-from json_stream.select_tokenizer import default_tokenizer
+from json_stream.select_tokenizer import get_tokenizer
 
 from json_stream.tests import JSONLoadTestCase
 
 
 @skipUnless(hasattr(json_stream_rs_tokenizer, 'RustTokenizer'), 'rust tokenizer not available')
 class TestRSTokenizer(JSONLoadTestCase):
     def test_load_object(self):
-        self.assertIs(default_tokenizer, json_stream_rs_tokenizer.RustTokenizer)
+        self.assertIs(get_tokenizer(), json_stream_rs_tokenizer.RustTokenizer)
         obj = {"a": 1, "b": None, "c": True}
         self._test_object(obj, persistent=False)
 
     def test_load_object_binary(self):
-        self.assertIs(default_tokenizer, json_stream_rs_tokenizer.RustTokenizer)
+        self.assertIs(get_tokenizer(), json_stream_rs_tokenizer.RustTokenizer)
         obj = {"a": 1, "b": None, "c": True}
         self._test_object(obj, persistent=False, binary=True)
 

src/json_stream/visitor.py

@@ -1,6 +1,5 @@
 from json_stream.base import StreamingJSONObject, StreamingJSONList, StreamingJSONBase
-from json_stream.iterators import ensure_file
-from json_stream.select_tokenizer import default_tokenizer
+from json_stream.select_tokenizer import get_token_stream
 
 
 def _visit(obj, visitor, path):
@@ -19,9 +18,8 @@ def _visit(obj, visitor, path):
         visitor(obj, path)
 
 
-def visit(fp_or_iterator, visitor, tokenizer=default_tokenizer, buffering=-1):
-    fp = ensure_file(fp_or_iterator)
-    token_stream = tokenizer(fp, buffering=buffering)
+def visit(fp_or_iterable, visitor, tokenizer=None, buffering=-1, **kwargs):
+    token_stream = get_token_stream(fp_or_iterable, tokenizer=tokenizer, buffering=buffering, **kwargs)
     _, token = next(token_stream)
     obj = StreamingJSONBase.factory(token, token_stream, persistent=False)
     _visit(obj, visitor, ())