Fixed -* list args, bash completion, zsh balloon completion

3 files changed, 38 insertions, 20 deletions

diff --git a/completion/bash-completion.sh b/completion/bash-completion.sh
index 6a346f5..f5284cf 100755
--- a/completion/bash-completion.sh
+++ b/completion/bash-completion.sh
@@ -1,21 +1,35 @@
-#!bash
-#
-# this is bash completion function file for ponysay.py.
-# generated by genzshcomp(ver: 0.5.1)
-#
+#!/usr/bin/env bash
 
-_ponysay.py()
+_ponysay()
 {
-  local cur
-  local cmd
+  local cur prev
+  cur=${COMP_WORDS[COMP_CWORD]}
+  prev=${COMP_WORDS[COMP_CWORD-1]}
 
-  cur=${COMP_WORDS[$COMP_CWORD]}
-  cmd=( ${COMP_WORDS[@]} )
+  declare -ga _ponysay_ponies _ponysay_balloons _ponythink_balloons
+  [[ 0 < "${#_ponysay_ponies}" ]] || _ponysay_ponies=( $(ponysay -p list) )
+  [[ 0 < "${#_ponysay_balloons}" ]] || _ponysay_balloons=( $(ponysay -b list) )
+  [[ 0 < "${#_ponythink_balloons}" ]] || _ponythink_balloons=( $(ponythink -b list) )
 
-  if [[ "$cur" == -* ]]; then
-    COMPREPLY=( $( compgen -W "-h --help --pony -p --quote -q --center -c --center-text -C --width -w --balloon -b" -- $cur ) )
-    return 0
+  if [[ "$cur" = -* ]]
+  then COMPREPLY=( $( compgen -W "-h --help --pony -p --quote -q --center -c --center-text -C --width -w --balloon -b" -- "$cur" ) )
+	return 0
+  fi
+
+  if [[ "$prev" =~ -[^-]*p$ || "$prev" = "--pony" ]]
+  then COMPREPLY=( $( compgen -W "${_ponysay_ponies[*]}" -- "$cur" ) )
+	return 0
+  fi
+
+  if [[ "$prev" =~ -[^-]*b$ || "$prev" = "--balloon" ]]
+  then
+	if [[ "$COMP_LINE" == ponythink* ]]
+	then COMPREPLY=( $( compgen -W "${_ponythink_balloons[*]}" -- "$cur" ) )
+	else COMPREPLY=( $( compgen -W "${_ponysay_balloons[*]}" -- "$cur" ) )
+	fi
+	return 0
   fi
 }
 
-complete -F _ponysay.py -o default ponysay.py
+complete -F _ponysay -o default ponysay
+complete -F _ponysay -o default ponythink
diff --git a/completion/zsh-completion.sh b/completion/zsh-completion.sh
index 4e280ba..917c918 100644
--- a/completion/zsh-completion.sh
+++ b/completion/zsh-completion.sh
@@ -14,8 +14,11 @@ local -a _available_ponies
   compadd "$@" -a _available_ponies
 }
 
-_ponysay_balloons(){
-  [[ "$words" == ponythink* ]] && compadd "$@" cowsay ascii || compadd "$@" cowsay ascii unicode round linux-vt
+local -a _ponysay_balloons _ponythink_balloons
+(( $+functions[_ponysay_pony] )) || _ponysay_balloons(){
+  (( $#_ponysay_balloons )) || _ponysay_balloons=($(ponysay -b list))
+  (( $#_ponythink_balloons )) || _ponythink_balloons=($(ponythink -b list))
+  [[ "$words" == ponythink* ]] && compadd "$@" -a _ponythink_balloons || compadd "$@" -a _ponysay_balloons
 }
 
 _arguments -s -w -S \
diff --git a/ponysay.py b/ponysay.py
index 23239a5..6b68463 100755
--- a/ponysay.py
+++ b/ponysay.py
@@ -94,12 +94,13 @@ if __name__ == '__main__':
 	parser.add_argument('-b', '--balloon', type=str, default='cowsay', help='Balloon style to use. Use "-b list" to list available styles.')
 	parser.add_argument('text', type=str, nargs='*', help='The text to be placed in the speech bubble')
 	args = parser.parse_args()
-
+  
+	think = sys.argv[0].endswith('think')
 	if args.pony == "list":
-		print('\n'.join(sorted(list_ponies(True))))
+		print('\n'.join(sorted(list_ponies() if not args.quote else list_ponies_with_quotes())))
 		sys.exit()
 	if args.balloon == 'list':
-		print('\n'.join(balloonstyles.keys()))
+		print('\n'.join([ s.replace('.think', '') for s in balloonstyles.keys() if s.endswith('.think') == think ]))
 		sys.exit()
 	pony = args.pony
 	if pony == "random":
@@ -111,7 +112,7 @@ if __name__ == '__main__':
 		text = random_quote(pony)
 
 	balloonstyle = None
-	if sys.argv[0].endswith('think'):
+	if think:
 		balloonstyle = balloonstyles[args.balloon+'.think']
 	else:
 		balloonstyle = balloonstyles[args.balloon]