#!/usr/bin/python2.7
"""Tool for inferring types from Python programs.

'pytype' is a tool for generating pyi from Python programs.

Usage:
  pytype [flags] file.py
"""

import atexit
import cProfile
import csv
import logging
import os
import sys
import traceback

from pytype import config
from pytype import directors
from pytype import errors
from pytype import infer
from pytype import metrics
from pytype.pyc import pyc
from pytype.pytd import optimize
from pytype.pytd import pytd
from pytype.pytd import utils as pytd_utils
from pytype.pytd.parse import node
from pytype.pytd.parse import visitors


log = logging.getLogger(__name__)


def check_pyi(input_filename, output_filename, errorlog, options):
  with open(input_filename, "r") as fi:
    py_src = fi.read()
  if output_filename is not None:
    with open(output_filename, "r") as fi:
      pyi_src = fi.read()
  else:
    pyi_src = None
  infer.check_types(
      py_src,
      pyi_src,
      py_filename=input_filename,
      pytd_filename=output_filename,
      errorlog=errorlog,
      options=options,
      run_builtins=options.run_builtins,
      deep=not options.main_only,
      cache_unknowns=options.cache_unknowns)


def generate_pyi(input_filename, errorlog, options):
  """Run the inferencer on one file, producing output.

  Args:
    input_filename: name of the file to process
    errorlog: Where error messages go. Instance of errors.ErrorLog.
    options: config.Options object.

  Returns:
    The pyi AST.

  Raises:
    CompileError: If we couldn't parse the input file.
  """
  with open(input_filename, "r") as fi:
    src = fi.read()

  mod, builtins = infer.infer_types(
      src,
      errorlog=errorlog,
      options=options,
      filename=input_filename,
      run_builtins=options.run_builtins,
      deep=not options.main_only,
      solve_unknowns=not options.quick,
      maximum_depth=1 if options.quick else 3,
      cache_unknowns=options.cache_unknowns)
  mod.Visit(visitors.VerifyVisitor())
  mod = optimize.Optimize(mod,
                          builtins,
                          # TODO(kramm): Add FLAGs for these
                          lossy=False,
                          use_abcs=False,
                          max_union=7,
                          remove_mutable=False)
  log.info("=========== pyi optimized =============")
  mod = pytd_utils.CanonicalOrdering(mod, sort_signatures=True)
  log.info("\n%s", pytd.Print(mod))
  log.info("========================================")

  result = pytd.Print(mod)
  if not result.endswith("\n"):  # TODO(pludemann): fix this hack
    result += "\n"
  result_prefix = ""
  if options.quick:
    result_prefix += "# (generated with --quick)\n"
  if result_prefix:
    result = result_prefix + "\n" + result
  return result


def process_one_file(input_filename,
                     output_filename,
                     options,
                     print_errors=True,
                     errors_csv_file=None):
  """Check or generate a .pyi, according to options.

  Args:
    input_filename: name of the file to process
    output_filename: name of the file for writing the output. If this is None,
                     then the options are used to determine where to write the
                     output.
    options: config.Options object.
    print_errors: whether to print the error log. This does not suppress all
                  errors (e.g., syntax errors) but is intended to suppress
                  possibly spurious messages during the first pass if pytype is
                  doing two passes.
    errors_csv_file: a csv writer object to write errors into. If None,
                     errors are not written into a csv file.

  Returns:
    An error code (0 means no error).

  """
  errorlog = errors.ErrorLog()
  with open(input_filename, "rb") as fi:
    director = directors.Director(fi.read(), errorlog, input_filename,
                                  options.disable)
  errorlog.set_error_filter(director.should_report_error)
  try:
    if options.check:
      check_pyi(input_filename=input_filename,
                output_filename=output_filename,
                errorlog=errorlog,
                options=options)
    else:
      # The default pyi that is used in the event that generate_pyi fails
      result = "def __getattr__(name) -> Any: ...\n"
      result = generate_pyi(input_filename=input_filename,
                            errorlog=errorlog,
                            options=options)
  except pyc.CompileError as e:
    errorlog.python_compiler_error(input_filename, e.lineno, e.error)
  except Exception as e:  # pylint: disable=broad-except
    if options.nofail:
      log.warn("***Caught exception: %s", str(e), exc_info=True)
      if not options.check:
        result += (  # pytype: disable=name-error
            "# Caught error in pytype: " + str(e).replace("\n", "\n#")
            + "\n# " + "\n# ".join(traceback.format_exc().splitlines()))
    else:
      raise
  if not options.check:
    if output_filename == "-" or not output_filename:
      sys.stdout.write(result)
    else:
      log.info("write pyi %r => %r", input_filename, output_filename)
      with open(output_filename, "w") as fi:
        fi.write(result)
  if options.report_errors:
    if print_errors:
      if errors_csv_file:
        errorlog.print_to_csv_file(errors_csv_file)
        return 0  # Command is successful regardless of errors.
      else:
        errorlog.print_to_stderr()

    return 1 if errorlog.has_error() else 0  # exit code
  else:
    return 0


class _ProfileContext(object):
  """A context manager for optionally profiling code."""

  def __init__(self, output_path):
    """Initialize.

    Args:
      output_path: A pathname for the profiler output.  An empty string
          indicates that no profiling should be done.
    """
    self._output_path = output_path
    self._profile = cProfile.Profile() if self._output_path else None

  def __enter__(self):
    if self._profile:
      self._profile.enable()

  def __exit__(self, exc_type, exc_value, traceback):  # pylint: disable=redefined-outer-name
    if self._profile:
      self._profile.disable()
      self._profile.dump_stats(self._output_path)


def main(argv):
  try:
    options = config.Options(argv)
  except config.OptParseError as e:
    print >>sys.stderr, str(e)
    sys.exit(1)

  with _ProfileContext(options.profile):
    with metrics.MetricsContext(options.metrics):
      with metrics.StopWatch("total_time"):
        return _run_pytype(options)


def _run_pytype(options):
  """Run pytype with the given configuration options."""
  if not options.src_out:
    print >>sys.stderr, "Need at least one filename."
    sys.exit(1)

  # TODO(dbaum): Consider changing flag default and/or polarity.  This will
  # need to be coordinated with a change to pytype.bzl.
  if not options.check_preconditions:
    node.DisablePreconditions()

  # Do *not* apply os.path.abspath here because we could be in a symlink tree
  # and bad things happen if you go to relative directories.

  # If we're processing more than one file, we need to do two passes (if we
  # don't know what the dependencies are). To speed things up, separate out the
  # biggest file and only process it once.  So, sort by size of the input files:
  if len(options.src_out) > 1:
    options.src_out.sort(reverse=True, key=lambda s: os.path.getsize(s[0]))
    for input_filename, output_filename in options.src_out[1:]:
      log.info("Process [pre-pass] %s => %s", input_filename, output_filename)
      _ = process_one_file(input_filename, output_filename, options,
                           print_errors=False)

  if options.output_errors_csv:
    f = open(options.output_errors_csv, "wb")
    errors_csv_file = csv.writer(f, delimiter=",")
    atexit.register(f.close)
  else:
    errors_csv_file = None

  exit_status = 0
  for input_filename, output_filename in options.src_out:
    log.info("Process %s => %s", input_filename, output_filename)
    ret = process_one_file(input_filename,
                           output_filename,
                           options,
                           print_errors=True,
                           errors_csv_file=errors_csv_file)
    exit_status = ret or exit_status

  # Touch output file upon success.
  if options.touch and not exit_status:
    with open(options.touch, "a"):
      os.utime(options.touch, None)

  return exit_status


if __name__ == "__main__":
  sys.exit(main(sys.argv) or 0)
