From 80a5e600ed8c5feb9189cb49e3fb95ffbf70af89 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Virtus?= <nechtom@gmail.com>
Date: Thu, 25 Jan 2024 11:19:47 +0100
Subject: [PATCH] Add optional args parameter to shortcut functions

compile() supports optional args parameter. Add it to shortcut functions
as well.
---
 README.rst        |  1 +
 jq.pyx            | 16 ++++++++--------
 tests/jq_tests.py | 23 +++++++++++++++++++++++
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/README.rst b/README.rst
index 9e90d73..349ffad 100644
--- a/README.rst
+++ b/README.rst
@@ -195,6 +195,7 @@ Convenience functions are available to get the output for a program and input in
 
     assert jq.first(".[] + 1", [1, 2, 3]) == 2
     assert jq.first(".[] + 1", text="[1, 2, 3]") == 2
+    assert jq.first(".[] + $addend", [1, 2, 3], args={"addend": 1}) == 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]
diff --git a/jq.pyx b/jq.pyx
index 0f580c4..4b10cd3 100644
--- a/jq.pyx
+++ b/jq.pyx
@@ -394,23 +394,23 @@ cdef class _ResultIterator(object):
             raise StopIteration()
 
 
-def all(program, value=_NO_VALUE, text=_NO_VALUE):
-    return compile(program).input(value, text=text).all()
+def all(program, value=_NO_VALUE, text=_NO_VALUE, args=None):
+    return compile(program, args=args).input(value, text=text).all()
 
 
-def first(program, value=_NO_VALUE, text=_NO_VALUE):
-    return compile(program).input(value, text=text).first()
+def first(program, value=_NO_VALUE, text=_NO_VALUE, args=None):
+    return compile(program, args=args).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 iter(program, value=_NO_VALUE, text=_NO_VALUE, args=None):
+    return _iter(compile(program, args=args).input(value, text=text))
 
 
-def text(program, value=_NO_VALUE, text=_NO_VALUE):
-    return compile(program).input(value, text=text).text()
+def text(program, value=_NO_VALUE, text=_NO_VALUE, args=None):
+    return compile(program, args=args).input(value, text=text).text()
 
 
 # Support the 0.1.x API for backwards compatibility
diff --git a/tests/jq_tests.py b/tests/jq_tests.py
index c911ef5..7e22fc0 100644
--- a/tests/jq_tests.py
+++ b/tests/jq_tests.py
@@ -431,3 +431,26 @@ def test_iter_function_with_json_text_input_returns_all_output_element_in_iterat
         assert_equal(3, next(iterator))
         assert_equal(4, next(iterator))
         assert_equal("end", next(iterator, "end"))
+
+    def test_all_function_with_args_should_set_predefined_variables(self):
+        output = jq.all(".[] | . + $a + $b", [1, 2, 3], args={"a": 100, "b": 20})
+
+        assert_equal([121, 122, 123], output)
+
+    def test_first_function_with_args_should_set_predefined_variables(self):
+        output = jq.first('.m = "a\\($x)c"', {"n": 1}, args=dict(x="b", m="z"))
+
+        assert_equal({"m": "abc", "n": 1}, output)
+
+    def test_iter_function_with_args_should_set_predefined_variables(self):
+        iterator = jq.iter(". * $add | .[]", {"a": 1, "b": 2}, args={"add": {"b": 3, "c": 4}})
+
+        assert_equal(1, next(iterator))
+        assert_equal(3, next(iterator))
+        assert_equal(4, next(iterator))
+        assert_equal("end", next(iterator, "end"))
+
+    def test_text_function_with_args_should_set_predefined_variables(self):
+        output = jq.text(". + $x", "foo", args=dict(x="bar"))
+
+        assert_equal('"foobar"', output)