From e7471d1a36ab71abdab4d1f3cd8f246108b42eb6 Mon Sep 17 00:00:00 2001
From: Michael Williamson <mike@zwobble.org>
Date: Sat, 2 May 2020 11:12:18 +0100
Subject: Add convenience functions

---
 README.rst        | 10 ++++++++++
 jq.pyx            | 21 +++++++++++++++++++-
 tests/jq_tests.py | 57 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+), 1 deletion(-)

diff --git a/README.rst b/README.rst
index 1217c3c..a3af21d 100644
--- a/README.rst
+++ b/README.rst
@@ -110,6 +110,16 @@ Call ``iter()`` to get all of the output elements as an iterator:
     assert next(iterator, None) == 4
     assert next(iterator, None) == None
 
+Convenience functions are available to get the output for a program and input in one call:
+
+.. code-block:: python
+
+    assert jq.first(".[] + 1", [1, 2, 3]) == 2
+    assert jq.first(".[] + 1", text="[1, 2, 3]") == 2
+    assert jq.text(".[] + 1", [1, 2, 3]) == "2\n3\n4"
+    assert jq.all(".[] + 1", [1, 2, 3]) == [2, 3, 4]
+    assert list(jq.iter(".[] + 1", [1, 2, 3])) == [2, 3, 4]
+
 The original program string is available on a compiled program as the ``program_string`` attribute:
 
 .. code-block:: python
diff --git a/jq.pyx b/jq.pyx
index 229faa4..f1c6a48 100644
--- a/jq.pyx
+++ b/jq.pyx
@@ -174,7 +174,7 @@ cdef class _ProgramWithInput(object):
         return list(self)
 
     def first(self):
-        return next(iter(self))
+        return next(_iter(self))
 
 
 cdef class _ResultIterator(object):
@@ -238,6 +238,25 @@ cdef class _ResultIterator(object):
             raise StopIteration()
 
 
+def all(program, value=_NO_VALUE, text=_NO_VALUE):
+    return compile(program).input(value, text=text).all()
+
+
+def first(program, value=_NO_VALUE, text=_NO_VALUE):
+    return compile(program).input(value, text=text).first()
+
+
+_iter = iter
+
+
+def iter(program, value=_NO_VALUE, text=_NO_VALUE):
+    return _iter(compile(program).input(value, text=text))
+
+
+def text(program, value=_NO_VALUE, text=_NO_VALUE):
+    return compile(program).input(value, text=text).text()
+
+
 # Support the 0.1.x API for backwards compatibility
 def jq(object program):
     return compile(program)
diff --git a/tests/jq_tests.py b/tests/jq_tests.py
index 17dfc0c..27734e6 100644
--- a/tests/jq_tests.py
+++ b/tests/jq_tests.py
@@ -195,3 +195,60 @@ def repr_of_compile_result_is_compilation_string():
 def program_string_can_be_retrieved_from_program():
     program = jq.compile(".")
     assert_equal(".", program.program_string)
+
+
+@istest
+class ConvenienceFunctions(object):
+    @istest
+    def first_function_with_json_value_input_returns_first_output_element(self):
+        output = jq.first(".[] + 1", [1, 2, 3])
+
+        assert_equal(2, output)
+
+    @istest
+    def first_function_with_json_text_input_returns_first_output_element(self):
+        output = jq.first(".[] + 1", text="[1, 2, 3]")
+
+        assert_equal(2, output)
+
+    @istest
+    def text_function_with_json_value_input_returns_all_output_elements_as_string(self):
+        output = jq.text(".[] + 1", [1, 2, 3])
+
+        assert_equal("2\n3\n4", output)
+
+    @istest
+    def text_function_with_json_text_input_returns_all_output_elements_as_string(self):
+        output = jq.text(".[] + 1", text="[1, 2, 3]")
+
+        assert_equal("2\n3\n4", output)
+
+    @istest
+    def all_function_with_json_value_input_returns_all_output_elements_in_list(self):
+        output = jq.all(".[] + 1", [1, 2, 3])
+
+        assert_equal([2, 3, 4], output)
+
+    @istest
+    def all_function_with_json_text_input_returns_all_output_element_in_list(self):
+        output = jq.all(".[] + 1", text="[1, 2, 3]")
+
+        assert_equal([2, 3, 4], output)
+
+    @istest
+    def iter_function_with_json_value_input_returns_all_output_elements_in_iterator(self):
+        iterator = jq.iter(".[] + 1", [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 iter_function_with_json_text_input_returns_all_output_element_in_iterator(self):
+        iterator = jq.iter(".[] + 1", text="[1, 2, 3]")
+
+        assert_equal(2, next(iterator))
+        assert_equal(3, next(iterator))
+        assert_equal(4, next(iterator))
+        assert_equal("end", next(iterator, "end"))
-- 
cgit v1.2.3