from typing import List
import types
import re


def _parse_pytype_output(var_names: List[str],
                         lines: List[str]) -> types.SimpleNamespace:
  """Parses the inferred type of each variable from pytype output."""
  reveal_type_lines = [l for l in lines if '[reveal-type]' in l]
  assert len(reveal_type_lines) == len(var_names)

  types_dict = {}
  for var, line in zip(var_names, reveal_type_lines):
    match = re.search(r'File "[^"]*", line \d+, in [^:]*: '
                      r'(.*) \[reveal-type\]', line)
    if match is None:
      raise Exception(f"Couldn't parse type from line: {line}")
    t = match.group(1)
    # Simplifies e.g. `tensor_annotations.jax.Array0` to just `Array0`
    t = re.sub(r'tensor_annotations.[^.]*\.', '', t)
    types_dict[var] = t

  return types.SimpleNamespace(**types_dict)