from typing import Dict
from typing import Any
import json
import re


def format_params_shapes(params_shapes: Dict[str, Any]) -> str:
  """Formats a dictionary of parameter shapes into a string.

  Args:
    params_shapes: Dictionary of parameter shapes.

  Returns:
    String formatted result of those parameter shapes, which is nicely formatted
      by using JSON / indentation, but formatting short lists into one line.
  """

  # Typically, shape arrays are very verbose, so we want to re-format them to
  # fit on a single line. Do so if it wouldn't overflow.
  def re_compact_arrays(array_match) -> str:
    try:
      values = json.loads(array_match.group(0))
    except ValueError:
      return array_match.group(0)
    re_compacted = json.dumps(values)  # no indent parameter
    return re_compacted if len(re_compacted) < 80 else array_match.group(0)

  json_formatted = json.dumps(params_shapes, indent=2)
  return re.sub(r'\[[^\[\]]+\]', re_compact_arrays, json_formatted)