34415: cache list of all python modules

3 files changed, 75 insertions, 25 deletions

diff --git a/Completion/Unix/Command/_pydoc b/Completion/Unix/Command/_pydoc
index f85ab0c2d..677c96ad6 100644
--- a/Completion/Unix/Command/_pydoc
+++ b/Completion/Unix/Command/_pydoc
@@ -1,19 +1,34 @@
-#compdef pydoc
+#compdef -P pydoc[0-9.]#
 
-local context state line ret=1
+local curcontext=$curcontext state state_descr line ret=1
 typeset -A opt_args
+local -a args
 
-_arguments \
-  '(-)-k[search keyword]:keyword' \
-  '(- *)-p[start web server on specified port]:port number' \
-  '(- *)-g[start gui]' \
-  '(-)-w[write out HTML]:file or dir:_files' \
-  '(- *)-h[show help information]' \
-  '*: :->lookup' && ret=0
+args=(
+  '(- *)-k[search keyword]:keyword'
+  '(-k -g -w *)-p[start web server on specified port]:port number'
+  '(-)-w[write out HTML in current directory]'
+  '(-)*: :->lookup'
+)
 
-[[ -n $state ]] && _alternative -C $context \
-  'keywords:keyword:compadd ${=${${(f)"$(_call_program keywords pydoc keywords)"}[2,-1]}}' \
-  'topics:topic:compadd ${=${${(f)"$(_call_program topics pydoc topics)"}[2,-1]}}' \
-  'modules:module:' && ret=0
+if _pick_variant pydoc3='pydoc3 -b' pydoc2 -h; then
+  args+=( '(-k -w *)-b[start server and open browser]' )
+else
+  args+=( '(- *)-g[start gui]' )
+fi
+
+_arguments -C : $args && return 0
+
+case $state in
+(lookup)
+  if [[ $words[CURRENT] = */* ]]; then
+    _files && ret=0
+  else
+    _alternative \
+      'keywords:keyword:compadd ${=${${(f)"$(_call_program keywords $words[1] keywords)"}[2,-1]}}' \
+      'topics:topic:compadd ${=${${(f)"$(_call_program topics $words[1] topics)"}[2,-1]}}' \
+      'modules:module:_python_modules' && ret=0
+  fi
+esac
 
 return ret
diff --git a/Completion/Unix/Command/_python b/Completion/Unix/Command/_python
index 8e32bdf55..dedb9ceda 100644
--- a/Completion/Unix/Command/_python
+++ b/Completion/Unix/Command/_python
@@ -1,11 +1,10 @@
-#compdef python python2 python2.4 python2.5 python2.6 python2.7 python3.0 python3.1 python3.2 python3.3 python3.4
+#compdef -P python[0-9.]#
 
 # Python 2.7
 # Python 3.4
 
-local curcontext="$curcontext" state line expl
+local curcontext="$curcontext" state state_descr line
 typeset -A opt_args
-
 local -a args
 
 if _pick_variant python3=Python\ 3 python2 --version; then
@@ -32,7 +31,7 @@ _arguments -C -s -S "$args[@]" \
   '-E[ignore PYTHON* environment variables (such as PYTHONPATH)]' \
   '(1 * -)-h[display help information]' \
   '-i[inspect interactively after running script]' \
-  '(1 -)-m[run library module as a script (terminates option list)]:module:->modules' \
+  '(1 -)-m[run library module as a script (terminates option list)]:module:_python_modules' \
   '-O[optimize generated bytecode slightly]' \
   '-OO[remove doc-strings in addition to the -O optimizations]' \
   "-s[don't add user site directory to sys.path]" \
@@ -46,14 +45,6 @@ _arguments -C -s -S "$args[@]" \
   '*::script argument:= ->normal' && return
 
 case "$state" in
-  modules)
-    local -a modules
-    modules=(
-      ${${=${(f)"$(_call_program modules $words[1] -c \
-        'from\ pydoc\ import\ help\;\ help\(\"modules\"\)')"}[2,-3]}:#\(package\)}
-    )
-    _wanted modules expl module compadd -a modules && return
-    ;;
   normal)
     if [[ -z "$opt_args[(I)-(c|m)]" ]]; then
       shift words
diff --git a/Completion/Unix/Type/_python_modules b/Completion/Unix/Type/_python_modules
new file mode 100644
index 000000000..b30848d83
--- /dev/null
+++ b/Completion/Unix/Type/_python_modules
@@ -0,0 +1,44 @@
+#autoload
+
+_python_module_caching_policy () {
+  local -a newer
+  # rebuild if cache does not exist or is more than a week old
+  newer=( "$1"(Nmw-1) )
+  return $#newer
+}
+
+_python_modules () {
+  local update_policy python expl
+
+  case $words[1] in
+    (python*) python=$words[1] ;;
+    (pydoc*)  python=${words[1]/#pydoc/python} ;;
+    (*)       python="python" ;;
+  esac
+  local cache_id=${${python//[^[:alnum:]]/_}#_}_modules
+  local array_name=_${cache_id}
+
+  zstyle -s ":completion:${curcontext}:" cache-policy update_policy
+  [[ -z "$update_policy" ]] && \
+    zstyle ":completion:${curcontext}:" \
+      cache-policy _python_module_caching_policy
+
+  if ( [[ ${(P)+array_name} -eq 0 ]] || _cache_invalid $cache_id ) &&
+      ! _retrieve_cache $cache_id; then
+
+    local script='import sys, pydoc
+def f(p,m,d):
+    if m.find(".") < 0: sys.stdout.write(m+"\n")
+pydoc.ModuleScanner().run(f)'
+
+    typeset -agU $array_name
+    set -A $array_name \
+      $(_call_program modules $python -c ${(q)script} 2>/dev/null)
+
+    _store_cache $cache_id $array_name
+  fi
+
+  _wanted modules expl module compadd "$@" -a -- $array_name
+}
+
+_python_modules "$@"