#compdef -subscript-

if [[ ${(Pt)${compstate[parameter]}} = assoc* ]]; then
  compgen -S ']' -k "( ${(kP)${compstate[parameter]}} )"
else
  _compalso -math-
fi