Design new API

1 files changed, 98 insertions, 65 deletions

diff --git a/jq.pyx b/jq.pyx
index 1bd443e..a27e238 100644
--- a/jq.pyx
+++ b/jq.pyx
@@ -1,5 +1,7 @@
 import json
 
+from cpython.bytes cimport PyBytes_AsString
+
 
 cdef extern from "jv.h":
     ctypedef enum jv_kind:
@@ -48,17 +50,24 @@ cdef extern from "jq.h":
     void jq_get_error_cb(jq_state *, jq_err_cb *, void **)
 
 
-def jq(object program):
-    cdef object program_bytes_obj = program.encode("utf8")
-    cdef char* program_bytes = program_bytes_obj
-    cdef jq_state *jq = jq_init()
-    if not jq:
-        raise Exception("jq_init failed")
+def compile(object program):
+    cdef object program_bytes = program.encode("utf8")
+    cdef jq_state* state = _compile(program_bytes)
+    jq_teardown(&state)
+
+    cdef _Program wrapped_program = _Program.__new__(_Program)
+    wrapped_program._program_bytes = program_bytes
+    return wrapped_program
 
+
+cdef jq_state* _compile(object program_bytes) except NULL:
+    cdef jq_state *jq = jq_init()
     cdef _ErrorStore error_store
     cdef int compiled
-
     try:
+        if not jq:
+            raise Exception("jq_init failed")
+
         error_store = _ErrorStore.__new__(_ErrorStore)
         error_store.clear()
 
@@ -71,22 +80,19 @@ def jq(object program):
 
         if not compiled:
             raise ValueError("program was not valid")
-
     except:
         jq_teardown(&jq)
         raise
+    # TODO: unset error callback?
 
-    cdef _Program wrapped_program = _Program.__new__(_Program)
-    wrapped_program._jq = jq
-    wrapped_program._error_store = error_store
-    return wrapped_program
+    return jq
 
 
 cdef void store_error(void* store_ptr, jv error):
     # TODO: handle errors not of JV_KIND_STRING
     cdef _ErrorStore store = <_ErrorStore>store_ptr
     if jv_get_kind(error) == JV_KIND_STRING:
-        store.store_error(jv_string_value(error))
+        store.store_error(jv_string_value(error).decode("utf8"))
 
     jv_free(error)
 
@@ -100,8 +106,8 @@ cdef class _ErrorStore(object):
     cdef object error_string(self):
         return "\n".join(self._errors)
 
-    cdef void store_error(self, char* error):
-        self._errors.append(error.decode("utf8"))
+    cdef void store_error(self, unicode error):
+        self._errors.append(error)
 
     cdef void clear(self):
         self._errors = []
@@ -113,76 +119,103 @@ class EmptyValue(object):
 _NO_VALUE = EmptyValue()
 
 cdef class _Program(object):
-    cdef jq_state* _jq
-    cdef _ErrorStore _error_store
-
-    def __dealloc__(self):
-        jq_teardown(&self._jq)
+    cdef object _program_bytes
 
-    def transform(self, value=_NO_VALUE, text=_NO_VALUE, text_output=False, multiple_output=False):
+    def input(self, value=_NO_VALUE, text=_NO_VALUE):
         if (value is _NO_VALUE) == (text is _NO_VALUE):
             raise ValueError("Either the value or text argument should be set")
         string_input = text if text is not _NO_VALUE else json.dumps(value)
-        bytes_input = string_input.encode("utf8")
 
-        self._error_store.clear()
+        cdef _ProgramWithInput program_with_input = _ProgramWithInput.__new__(_ProgramWithInput)
+        program_with_input._program_bytes = self._program_bytes
+        program_with_input._bytes_input = string_input.encode("utf8")
+        return program_with_input
 
-        result_bytes = self._string_to_strings(bytes_input)
 
-        if self._error_store.has_errors():
-            raise ValueError(self._error_store.error_string())
+cdef class _ProgramWithInput(object):
+    cdef object _program_bytes
+    cdef object _bytes_input
 
-        result_strings = map(lambda s: s.decode("utf8"), result_bytes)
-        if text_output:
-            return "\n".join(result_strings)
-        elif multiple_output:
-            return [json.loads(s) for s in result_strings]
-        else:
-            return json.loads(next(iter(result_strings)))
+    def __iter__(self):
+        return self._make_iterator()
 
-    cdef object _string_to_strings(self, char* input):
-        cdef jv_parser* parser = jv_parser_new(0)
-        jv_parser_set_buf(parser, input, len(input), 0)
-        cdef jv value
-        cdef jv error_message
+    cdef _ResultIterator _make_iterator(self):
+        iterator = _ResultIterator.__new__(_ResultIterator)
+        iterator._init(self._program_bytes, self._bytes_input)
+        return iterator
+
+    def text(self):
+        iterator = self._make_iterator()
         results = []
         while True:
-            value = jv_parser_next(parser)
-            if jv_is_valid(value):
-                self._process(value, results)
-            else:
-                self._handle_invalid_jv(value, b"parse error: ")
-                break
+            try:
+                results.append(iterator._next_string())
+            except StopIteration:
+                return "\n".join(results)
 
-        jv_parser_free(parser)
+    def all(self):
+        return list(self)
 
-        return results
+    def first(self):
+        return next(iter(self))
 
-    cdef void _handle_invalid_jv(self, jv value, char* prefix):
-        if jv_invalid_has_msg(jv_copy(value)):
-            error_message = jv_invalid_get_msg(value)
-            full_error_message = prefix + jv_string_value(error_message)
-            self._error_store.store_error(full_error_message)
-            jv_free(error_message)
-        else:
-            jv_free(value)
 
+cdef class _ResultIterator(object):
+    cdef jq_state* _jq
+    cdef jv_parser* _parser
+    cdef object _bytes_input
+    cdef bint _ready
+
+    def __dealloc__(self):
+        jq_teardown(&self._jq)
+        jv_parser_free(self._parser)
 
+    def _init(self, object program_bytes, object bytes_input):
+        self._jq = _compile(program_bytes)
+        self._bytes_input = bytes_input
+        self._ready = False
+        cdef jv_parser* parser = jv_parser_new(0)
+        cdef char* cbytes_input = PyBytes_AsString(bytes_input)
+        jv_parser_set_buf(parser, cbytes_input, len(cbytes_input), 0)
+        self._parser = parser
 
-    cdef void _process(self, jv value, object output):
-        cdef int jq_flags = 0
+    def __iter__(self):
+        return self
 
-        jq_start(self._jq, value, jq_flags);
-        cdef jv result
-        cdef int dumpopts = 0
-        cdef jv dumped
+    def __next__(self):
+        return json.loads(self._next_string())
 
+    cdef unicode _next_string(self):
+        cdef int dumpopts = 0
         while True:
+            if not self._ready:
+                self._ready_next_input()
+                self._ready = True
+
             result = jq_next(self._jq)
-            if not jv_is_valid(result):
-                self._handle_invalid_jv(result, b"")
-                return
-            else:
+            if jv_is_valid(result):
                 dumped = jv_dump_string(result, dumpopts)
-                output.append(jv_string_value(dumped))
+                value = jv_string_value(dumped).decode("utf8")
                 jv_free(dumped)
+                return value
+            elif jv_invalid_has_msg(jv_copy(result)):
+                error_message = jv_invalid_get_msg(result)
+                message = jv_string_value(error_message).decode("utf8")
+                jv_free(error_message)
+                raise ValueError(message)
+            else:
+                self._ready = False
+
+    cdef bint _ready_next_input(self) except 1:
+        cdef int jq_flags = 0
+        cdef jv value = jv_parser_next(self._parser)
+        if jv_is_valid(value):
+            jq_start(self._jq, value, jq_flags)
+            return 0
+        elif jv_invalid_has_msg(jv_copy(value)):
+            error_message = jv_invalid_get_msg(value)
+            message = jv_string_value(error_message).decode("utf8")
+            jv_free(error_message)
+            raise ValueError(u"parse error: " + message)
+        else:
+            raise StopIteration()