Design new API

3 files changed, 185 insertions, 119 deletions

diff --git a/README.rst b/README.rst
index 18aae1c..6eb1f65 100644
--- a/README.rst
+++ b/README.rst
@@ -60,50 +60,48 @@ Once Homebrew is installed, you can install the remaining dependencies with:
 Usage
 -----
 
-A program can be compiled by passing it to ``jq.jq``.
-To apply the program to an input, call the ``transform`` method.
-jq.py expects the value to be valid JSON,
-such as values returned from ``json.load``.
+Call ``jq.compile`` to compile a jq program.
+Call ``.input()`` on the compiled program to supply an input value.
+The input must either be:
+  * a valid JSON value, such as the values returned from ``json.load``
+  * unparsed JSON text passed as the keyword argument ``text``.
 
-.. code-block:: python
-
-    from jq import jq
-
-    jq(".").transform("42") == "42"
-    jq(".").transform({"a": 1}) == {"a": 1}
-
-If the value is unparsed JSON text, pass it in using the ``text``
-argument:
+Calling ``first()`` on the result will run the program with the given input,
+and return the first output element.
 
 .. code-block:: python
 
-    jq(".").transform(text="42") == 42
+    import jq
+
+    assert jq.compile(".").input("hello").first() == "hello"
+    assert jq.compile(".").input(text='"hello"').first() == "hello"
+    assert jq.compile("[.[]+1]").input([1, 2, 3]).first() == [2, 3, 4]
+    assert jq.compile(".[]+1").input([1, 2, 3]).first() == 2
 
-The ``text_output`` argument can be used to serialise the output into
-JSON text:
+Call ``text()`` instead of ``first()`` to serialise the output into JSON text:
 
 .. code-block:: python
 
-    jq(".").transform("42", text_output=True) == '"42"'
+    assert jq.compile(".").input("42").text() == '"42"'
 
-If there are multiple output elements, each element is represented by a
-separate line, irrespective of the value of ``multiple_output``:
+When calling ``text()``, if there are multiple output elements, each element is represented by a separate line:
 
 .. code-block:: python
 
-    jq(".[]").transform([1, 2, 3], text_output=True) == "1\n2\n3"
+    assert jq.compile(".[]").input([1, 2, 3]).text() == "1\n2\n3"
 
-If ``multiple_output`` is ``False`` (the default), then the first output
-is used:
+Call ``all()`` to get all of the output elements in a list:
 
 .. code-block:: python
 
-    jq(".[]+1").transform([1, 2, 3]) == 2
+    assert jq.compile(".[]+1").input([1, 2, 3]).all() == [2, 3, 4]
 
-If ``multiple_output`` is ``True``, all output elements are returned in
-an array:
+Call ``iter()`` to get all of the output elements as an iterator:
 
 .. code-block:: python
 
-    jq(".[]+1").transform([1, 2, 3], multiple_output=True) == [2, 3, 4]
-
+    iterator = iter(jq.compile(".[]+1").input([1, 2, 3]))
+    assert next(iterator, None) == 2
+    assert next(iterator, None) == 3
+    assert next(iterator, None) == 4
+    assert next(iterator, None) == None
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()
diff --git a/tests/jq_tests.py b/tests/jq_tests.py
index 4c37591..b6c3c38 100644
--- a/tests/jq_tests.py
+++ b/tests/jq_tests.py
@@ -4,14 +4,14 @@ from __future__ import unicode_literals
 
 from nose.tools import istest, assert_equal, assert_raises
 
-from jq import jq
+import jq
 
 
 @istest
 def output_of_dot_operator_is_input():
     assert_equal(
         "42",
-        jq(".").transform("42")
+        jq.compile(".").input("42").first()
     )
 
 
@@ -19,7 +19,7 @@ def output_of_dot_operator_is_input():
 def can_add_one_to_each_element_of_an_array():
     assert_equal(
         [2, 3, 4],
-        jq("[.[]+1]").transform([1, 2, 3])
+        jq.compile("[.[]+1]").input([1, 2, 3]).first()
     )
 
 
@@ -27,33 +27,33 @@ def can_add_one_to_each_element_of_an_array():
 def can_use_regexes():
     assert_equal(
         True,
-        jq('test(".*")').transform("42")
+        jq.compile('test(".*")').input("42").first()
     )
 
     assert_equal(
         True,
-        jq('test("^[0-9]+$")').transform("42")
+        jq.compile('test("^[0-9]+$")').input("42").first()
     )
 
     assert_equal(
         False,
-        jq('test("^[0-9]+$")').transform("42a")
+        jq.compile('test("^[0-9]+$")').input("42a").first()
     )
 
 
 @istest
-def input_string_is_parsed_to_json_if_raw_input_is_true():
+def when_text_argument_is_used_then_input_is_treated_as_json_text():
     assert_equal(
         42,
-        jq(".").transform(text="42")
+        jq.compile(".").input(text="42").first()
     )
 
 
 @istest
-def output_is_serialised_to_json_string_if_text_output_is_true():
+def when_text_method_is_used_on_result_then_output_is_serialised_to_json_string():
     assert_equal(
         '"42"',
-        jq(".").transform("42", text_output=True)
+        jq.compile(".").input("42").text()
     )
 
 
@@ -61,38 +61,73 @@ def output_is_serialised_to_json_string_if_text_output_is_true():
 def elements_in_text_output_are_separated_by_newlines():
     assert_equal(
         "1\n2\n3",
-        jq(".[]").transform([1, 2, 3], text_output=True)
+        jq.compile(".[]").input([1, 2, 3]).text()
     )
 
 
 @istest
-def first_output_element_is_returned_if_multiple_output_is_false_but_there_are_multiple_output_elements():
+def when_first_method_is_used_on_result_then_first_element_of_result_is_returned():
     assert_equal(
         2,
-        jq(".[]+1").transform([1, 2, 3])
+        jq.compile(".[]+1").input([1, 2, 3]).first()
     )
 
 
 @istest
-def multiple_output_elements_are_returned_if_multiple_output_is_true():
+def when_all_method_is_used_on_result_then_all_elements_are_returned_in_list():
     assert_equal(
         [2, 3, 4],
-        jq(".[]+1").transform([1, 2, 3], multiple_output=True)
+        jq.compile(".[]+1").input([1, 2, 3]).all()
     )
 
 
 @istest
-def multiple_inputs_in_raw_input_are_separated_by_newlines():
+def can_treat_execute_result_as_iterable():
+    iterator = iter(jq.compile(".[]+1").input([1, 2, 3]))
+    assert_equal(2, next(iterator))
+    assert_equal(3, next(iterator))
+    assert_equal(4, next(iterator))
+    assert_equal("end", next(iterator, "end"))
+
+
+@istest
+def can_execute_same_program_again_before_consuming_output_of_first_execution():
+    program = jq.compile(".[]+1")
+    first = iter(program.input([1, 2, 3]))
+    assert_equal(2, next(first))
+    second = iter(program.input([11, 12, 13]))
+    assert_equal(12, next(second))
+    assert_equal(3, next(first))
+    assert_equal(4, next(first))
+    assert_equal(13, next(second))
+    assert_equal(14, next(second))
+
+
+@istest
+def iterators_from_same_program_and_input_are_independent():
+    program_with_input = jq.compile(".[]+1").input([1, 2, 3])
+    first = iter(program_with_input)
+    assert_equal(2, next(first))
+    second = iter(program_with_input)
+    assert_equal(2, next(second))
+    assert_equal(3, next(first))
+    assert_equal(4, next(first))
+    assert_equal(3, next(second))
+    assert_equal(4, next(second))
+
+
+@istest
+def multiple_inputs_in_text_input_are_separated_by_newlines():
     assert_equal(
         [2, 3, 4],
-        jq(".+1").transform(text="1\n2\n3", multiple_output=True)
+        jq.compile(".+1").input(text="1\n2\n3").all()
     )
 
 
 @istest
 def value_error_is_raised_if_program_is_invalid():
     try:
-        jq("!")
+        jq.compile("!")
         assert False, "Expected error"
     except ValueError as error:
         expected_error_str = "jq: error: syntax error, unexpected INVALID_CHARACTER, expecting $end (Unix shell quoting issues?) at <top-level>, line 1:\n!\njq: 1 compile error"
@@ -101,9 +136,9 @@ def value_error_is_raised_if_program_is_invalid():
 
 @istest
 def value_error_is_raised_if_input_cannot_be_processed_by_program():
-    program = jq(".x")
+    program = jq.compile(".x")
     try:
-        program.transform(1)
+        program.input(1).all()
         assert False, "Expected error"
     except ValueError as error:
         expected_error_str = "Cannot index number with string \"x\""
@@ -112,21 +147,21 @@ def value_error_is_raised_if_input_cannot_be_processed_by_program():
 
 @istest
 def errors_do_not_leak_between_transformations():
-    program = jq(".x")
+    program = jq.compile(".x")
     try:
-        program.transform(1)
+        program.input(1).all()
         assert False, "Expected error"
     except ValueError as error:
         pass
 
-    assert_equal(1, program.transform({"x": 1}))
+    assert_equal(1, program.input({"x": 1}).first())
 
 
 @istest
 def value_error_is_raised_if_input_is_not_valid_json():
-    program = jq(".x")
+    program = jq.compile(".x")
     try:
-        program.transform(text="!!")
+        program.input(text="!!").first()
         assert False, "Expected error"
     except ValueError as error:
         expected_error_str = "parse error: Invalid numeric literal at EOF at line 1, column 2"
@@ -137,7 +172,7 @@ def value_error_is_raised_if_input_is_not_valid_json():
 def unicode_strings_can_be_used_as_input():
     assert_equal(
         "‽",
-        jq(".").transform(text='"‽"')
+        jq.compile(".").input(text='"‽"').first()
     )
 
 
@@ -145,5 +180,5 @@ def unicode_strings_can_be_used_as_input():
 def unicode_strings_can_be_used_as_programs():
     assert_equal(
         "Dragon‽",
-        jq('.+"‽"').transform(text='"Dragon"')
+        jq.compile('.+"‽"').input(text='"Dragon"').first()
     )