From dcf244279c13373cc344a77c4cc287c51fc65f45 Mon Sep 17 00:00:00 2001
From: Victor Hallberg <victor@hallberg.cc>
Date: Sun, 17 May 2020 23:47:48 +0200
Subject: [PATCH] Add g:jsonpath_use_python option which enables python
 implementation

---
 README.md             | 10 ++++++++
 autoload/jsonpath.vim | 53 +++++++++++++++++++++++++++++++++++++++----
 doc/jsonpath.txt      | 17 ++++++++++++++
 plugin/jsonpath.vim   |  4 ++++
 4 files changed, 79 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 10ada89..510a8d4 100644
--- a/README.md
+++ b/README.md
@@ -22,6 +22,16 @@ au FileType json noremap <buffer> <silent> <expr> <leader>p jsonpath#echo()
 au FileType json noremap <buffer> <silent> <expr> <leader>g jsonpath#goto()
 ```
 
+### Python support in vim
+
+While not required it is recommended to use a Vim environment with the
+`+python3` feature enabled, since the plugin provides a python implementation
+that is much faster than the vimscript variant. You can check the availability
+using:
+```vim
+:echo has("python3")
+```
+
 ## Mappings
 
 Mappings are not provided by default but can easily be added to your `.vimrc`.
diff --git a/autoload/jsonpath.vim b/autoload/jsonpath.vim
index 8ad19f6..0626960 100644
--- a/autoload/jsonpath.vim
+++ b/autoload/jsonpath.vim
@@ -1,11 +1,13 @@
 " autoload/jsonpath.vim
 " Author: Victor Hallberg <https://hallberg.cc>
 
-if exists("g:autoloaded_jsonpath")
+if exists('g:autoloaded_jsonpath')
   finish
 endif
 let g:autoloaded_jsonpath = 1
 
+let s:plugin_dir = expand('<sfile>:p:h:h')
+
 let s:escapes = {
   \ 'b': "\b",
   \ 'f': "\f",
@@ -45,12 +47,27 @@ function! jsonpath#scan_buffer(search_for, ...) "{{{
   endif
   let is_searching = !empty(search_for)
 
+  let to_column = max([0, get(a:, 2)])
+
   let to_line = get(a:, 1)
   if to_line < 1 || to_line > line('$')
     let to_line = line('$')
+    let to_column = strchars(getline('$'))
   endif
 
-  let to_column = max([0, get(a:, 2)])
+  if g:jsonpath_use_python
+    if has('python3')
+      return jsonpath#scan_buffer_python(search_for, to_line, to_column)
+    endif
+
+    echom 'g:jsonpath_use_python set but python not found, falling back to vimscript'
+  endif
+
+  return jsonpath#scan_buffer_vimscript(search_for, to_line, to_column)
+endfunction "}}}
+
+function! jsonpath#scan_buffer_vimscript(search_for, to_line, to_column) "{{{
+  let is_searching = !empty(a:search_for)
 
   " Parser state
   let stack = []
@@ -63,7 +80,7 @@ function! jsonpath#scan_buffer(search_for, ...) "{{{
 
   try
     let lnr = 1
-    while lnr <= to_line "{{{
+    while lnr <= a:to_line "{{{
       let line = getline(lnr)
       let line_length = len(line)
       let cnr = 1
@@ -129,13 +146,13 @@ function! jsonpath#scan_buffer(search_for, ...) "{{{
                 \})
           
           " Check if the sought search_for path has been reached?
-          if stack_modified == 1 && is_searching && s:is_equal_lists(stack, search_for)
+          if stack_modified == 1 && is_searching && s:is_equal_lists(stack, a:search_for)
             return [bufnr('%'), lnr, cnr, 0]
           endif
         endif
 
         " Abort if end position has been reached
-        if !parsing_key && lnr >= to_line && cnr + 1 >= to_column
+        if !parsing_key && lnr >= a:to_line && cnr + 1 >= a:to_column
           let finished = !is_searching " search failed if we reached end
           break
         endif
@@ -185,6 +202,32 @@ function! jsonpath#scan_buffer(search_for, ...) "{{{
   return []
 endfunction "}}}
 
+function! jsonpath#scan_buffer_python(search_for, to_line, to_column) "{{{
+py3 << EOF
+import sys
+import vim
+sys.path.insert(0, vim.eval('s:plugin_dir'))
+import jsonpath
+
+stream = jsonpath.CountingLines(vim.current.buffer)
+result = jsonpath.scan_stream(
+  stream,
+  path=vim.eval('a:search_for'),
+  line=int(vim.eval('a:to_line')),
+  column=int(vim.eval('a:to_column')),
+)
+EOF
+
+  let result = py3eval('result')
+  if empty(result)
+    return []
+  elseif !empty(a:search_for)
+    return [bufnr('%'), result[0], result[1], 0]
+  endif
+
+  return result
+endfunction "}}}
+
 " Attempts to place the cursor on identifier for the given path
 function! jsonpath#goto(...) "{{{
   let search_for = get(a:, 1)
diff --git a/doc/jsonpath.txt b/doc/jsonpath.txt
index 8f869f4..8f13db5 100644
--- a/doc/jsonpath.txt
+++ b/doc/jsonpath.txt
@@ -11,6 +11,12 @@ INTRODUCTION                                                        *jsonpath*
 
 This plugin provides ways of navigating JSON document buffers.
 
+While not required it is recommended to use a Vim environment with the
+|+python3| feature enabled, since the python implementation is much faster
+than the vimscript variant. You can check the availability using: >
+    :echo has("python3")
+<
+
                                                        *jsonpath-dot-notation*
 Paths should be entered (and will be output) using dot notation, consisting of
 properties separated by dots (configurable, see |g:jsonpath_delimeter|).
@@ -63,6 +69,17 @@ Set using >
     let g:jsonpath_register = '*'
 <
 
+|g:jsonpath_use_python|                                *g:jsonpath_use_python*
+Type: |int|
+Default: `has('python3')`
+
+Set this value to `0` to prevent the plugin from using python to scan the
+buffer. This option should generally be left alone as the python
+implementation is much faster.
+Override using >
+    let g:jsonpath_use_python = 0
+<
+
 ==============================================================================
 COMMANDS                                                   *jsonpath-commands*
 
diff --git a/plugin/jsonpath.vim b/plugin/jsonpath.vim
index 040f7c2..7a488a3 100644
--- a/plugin/jsonpath.vim
+++ b/plugin/jsonpath.vim
@@ -10,6 +10,10 @@ if !exists('g:jsonpath_delimeter')
   let g:jsonpath_delimeter = '.'
 endif
 
+if !exists('g:jsonpath_use_python')
+  let g:jsonpath_use_python = has('python3')
+endif
+
 command! -nargs=? JsonPath call jsonpath#command(<q-args>)
 
 " au FileType json noremap <buffer> <silent> <expr> <leader>g jsonpath#goto()