#compdef -subscript-

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