Add support for predefined variables

3 files changed, 35 insertions, 9 deletions

diff --git a/README.rst b/README.rst
index d331181..affe1f2 100644
--- a/README.rst
+++ b/README.rst
@@ -106,6 +106,13 @@ Call ``iter()`` to get all of the output elements as an iterator:
     assert next(iterator, None) == 4
     assert next(iterator, None) == None
 
+Calling ``compile()`` with the ``args`` argument allows predefined variables to be used within the program:
+
+.. code-block:: python
+
+    program = jq.compile("$a + $b + .", args={"a": 100, "b": 20})
+    assert program.input(3).first() == 123
+
 Convenience functions are available to get the output for a program and input in one call:
 
 .. code-block:: python
diff --git a/jq.pyx b/jq.pyx
index 7577ae5..d3f9067 100644
--- a/jq.pyx
+++ b/jq.pyx
@@ -44,6 +44,8 @@ cdef extern from "jv.h":
     void jv_parser_set_buf(jv_parser*, const char*, int, int)
     jv jv_parser_next(jv_parser*)
 
+    jv jv_parse(const char*)
+
 
 cdef extern from "jq.h":
     ctypedef struct jq_state:
@@ -54,6 +56,7 @@ cdef extern from "jq.h":
     jq_state *jq_init()
     void jq_teardown(jq_state **)
     int jq_compile(jq_state *, const char* str)
+    int jq_compile_args(jq_state *, const char* str, jv)
     void jq_start(jq_state *, jv value, int flags)
     jv jq_next(jq_state *)
     void jq_set_error_cb(jq_state *, jq_err_cb, void *)
@@ -106,17 +109,18 @@ cdef object _jv_to_python(jv value):
     return python_value
 
 
-def compile(object program):
+def compile(object program, args=None):
     cdef object program_bytes = program.encode("utf8")
-    return _Program(program_bytes)
+    return _Program(program_bytes, args=args)
 
 
 _compilation_lock = threading.Lock()
 
 
-cdef jq_state* _compile(object program_bytes) except NULL:
+cdef jq_state* _compile(object program_bytes, object args) except NULL:
     cdef jq_state *jq = jq_init()
     cdef _ErrorStore error_store
+    cdef jv jv_args
     cdef int compiled
     try:
         if not jq:
@@ -127,7 +131,12 @@ cdef jq_state* _compile(object program_bytes) except NULL:
         with _compilation_lock:
             jq_set_error_cb(jq, _store_error, <void*>error_store)
 
-            compiled = jq_compile(jq, program_bytes)
+            if args is None:
+                compiled = jq_compile(jq, program_bytes)
+            else:
+                args_bytes = json.dumps(args).encode("utf-8")
+                jv_args = jv_parse(PyBytes_AsString(args_bytes))
+                compiled = jq_compile_args(jq, program_bytes, jv_args)
 
             if error_store.has_errors():
                 raise ValueError(error_store.error_string())
@@ -179,11 +188,13 @@ _NO_VALUE = _EmptyValue()
 cdef class _JqStatePool(object):
     cdef jq_state* _jq_state
     cdef object _program_bytes
+    cdef object _args
     cdef object _lock
 
-    def __cinit__(self, program_bytes):
+    def __cinit__(self, program_bytes, args):
         self._program_bytes = program_bytes
-        self._jq_state = _compile(self._program_bytes)
+        self._args = args
+        self._jq_state = _compile(self._program_bytes, args=self._args)
         self._lock = threading.Lock()
 
     def __dealloc__(self):
@@ -192,7 +203,7 @@ cdef class _JqStatePool(object):
     cdef jq_state* acquire(self):
         with self._lock:
             if self._jq_state == NULL:
-                return _compile(self._program_bytes)
+                return _compile(self._program_bytes, args=self._args)
             else:
                 state = self._jq_state
                 self._jq_state = NULL
@@ -210,9 +221,9 @@ cdef class _Program(object):
     cdef object _program_bytes
     cdef _JqStatePool _jq_state_pool
 
-    def __cinit__(self, program_bytes):
+    def __cinit__(self, program_bytes, args):
         self._program_bytes = program_bytes
-        self._jq_state_pool = _JqStatePool(program_bytes)
+        self._jq_state_pool = _JqStatePool(program_bytes, args=args)
 
     def input(self, value=_NO_VALUE, text=_NO_VALUE):
         if (value is _NO_VALUE) == (text is _NO_VALUE):
diff --git a/tests/jq_tests.py b/tests/jq_tests.py
index 6b63556..7956d56 100644
--- a/tests/jq_tests.py
+++ b/tests/jq_tests.py
@@ -185,6 +185,14 @@ def unicode_strings_can_be_used_as_programs():
 
 
 @istest
+def compiling_with_args_should_set_predefined_variables():
+    assert_equal(
+        123,
+        jq.compile("$a + $b + .", args={"a": 100, "b": 20}).input(3).first()
+    )
+
+
+@istest
 def repr_of_compile_result_is_compilation_string():
     program = jq.compile(".")
     repr_string = repr(program)