# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Run doctests for tensorflow."""

import ast
import doctest
import os
import re
import textwrap
from typing import Any, Callable, Dict, Iterable, Optional

import astor

from tensorflow.tools.docs import tf_doctest_lib


def load_from_files(
    files,
    globs: Optional[Dict[str, Any]] = None,
    set_up: Optional[Callable[[Any], None]] = None,
    tear_down: Optional[Callable[[Any], None]] = None) -> doctest.DocFileSuite:
  """Creates a doctest suite from the files list.

  Args:
    files: A list of file paths to test.
    globs: The global namespace the tests are run in.
    set_up: Run before each test, receives the test as argument.
    tear_down: Run after each test, receives the test as argument.

  Returns:
    A DocFileSuite containing the tests.
  """
  if globs is None:
    globs = {}

  # __fspath__ isn't respected everywhere in doctest so convert paths to
  # strings.
  files = [os.fspath(f) for f in files]

  globs['_print_if_not_none'] = _print_if_not_none
  # Ref: https://docs.python.org/3/library/doctest.html#doctest.DocFileSuite
  return doctest.DocFileSuite(
      *files,
      module_relative=False,
      parser=FencedCellParser(fence_label='python'),
      globs=globs,
      setUp=set_up,
      tearDown=tear_down,
      checker=FencedCellOutputChecker(),
      optionflags=(doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE
                   | doctest.IGNORE_EXCEPTION_DETAIL
                   | doctest.DONT_ACCEPT_BLANKLINE),
  )


class FencedCellOutputChecker(tf_doctest_lib.TfDoctestOutputChecker):
  """TfDoctestChecker with a different warning message."""
  MESSAGE = textwrap.dedent("""\n
        ##############################################################
        # Check the documentation (go/g3doctest) on how to write
        # testable g3docs.
        ##############################################################
        """)


class FencedCellParser(doctest.DocTestParser):
  """Implements test parsing for ``` fenced cells.

  https://docs.python.org/3/library/doctest.html#doctestparser-objects

  The `get_examples` method receives a string and returns an
  iterable of `doctest.Example` objects.
  """
  patched = False

  def __init__(self, fence_label='python'):
    super().__init__()

    if not self.patched:
      # The default doctest compiles in "single" mode. The fenced block may
      # contain multiple statements. The `_patch_compile` function fixes the
      # compile mode.
      doctest.compile = _patch_compile
      print(
          textwrap.dedent("""
          *********************************************************************
          * Caution: `fenced_doctest` patches `doctest.compile` don't use this
          *   in the same binary as any other doctests.
          *********************************************************************
          """))
      type(self).patched = True

    # Match anything, except if the look-behind sees a closing fence.
    no_fence = '(.(?<!```))*?'
    self.fence_cell_re = re.compile(
        rf"""
        ^(                             # After a newline
            \s*```\s*({fence_label})\n   # Open a labeled ``` fence
            (?P<doctest>{no_fence})      # Match anything except a closing fence
            \n\s*```\s*(\n|$)            # Close the fence.
        )
        (                              # Optional!
            [\s\n]*                      # Any number of blank lines.
            ```\s*\n                     # Open ```
            (?P<output>{no_fence})       # Anything except a closing fence
            \n\s*```                     # Close the fence.
        )?
        """,
        # Multiline so ^ matches after a newline
        re.MULTILINE |
        # Dotall so `.` matches newlines.
        re.DOTALL |
        # Verbose to allow comments/ignore-whitespace.
        re.VERBOSE)

  def get_examples(self,
                   string: str,
                   name: str = '<string>') -> Iterable[doctest.Example]:
    # Check for a file-level skip comment.
    if re.search('<!--.*?doctest.*?skip.*?all.*?-->', string, re.IGNORECASE):
      return

    for match in self.fence_cell_re.finditer(string):
      if re.search('doctest.*skip', match.group(0), re.IGNORECASE):
        continue

      groups = match.groupdict()

      source = textwrap.dedent(groups['doctest'])
      want = groups['output']
      if want is not None:
        want = textwrap.dedent(want)

      yield doctest.Example(
          lineno=string[:match.start()].count('\n') + 1,
          source=source,
          want=want)


def _print_if_not_none(obj):
  """Print like a notebook: Show the repr if the object is not None.

  `_patch_compile` Uses this on the final expression in each cell.

  This way the outputs feel like notebooks.

  Args:
    obj: the object to print.
  """
  if obj is not None:
    print(repr(obj))


def _patch_compile(source,
                   filename,
                   mode,
                   flags=0,
                   dont_inherit=False,
                   optimize=-1):
  """Patch `doctest.compile` to make doctest to behave like a notebook.

  Default settings for doctest are configured to run like a repl: one statement
  at a time. The doctest source uses `compile(..., mode="single")`

  So to let doctest act like a notebook:

  1. We need `mode="exec"` (easy)
  2. We need the last expression to be printed (harder).

  To print the last expression, just wrap the last expression in
  `_print_if_not_none(expr)`. To detect the last expression use `AST`.
  If the last node is an expression modify the ast to call
  `_print_if_not_none` on it, convert the ast back to source and compile that.

  https://docs.python.org/3/library/functions.html#compile

  Args:
    source: Can either be a normal string, a byte string, or an AST object.
    filename: Argument should give the file from which the code was read; pass
      some recognizable value if it wasn’t read from a file ('<string>' is
      commonly used).
    mode: [Ignored] always use exec.
    flags: Compiler options.
    dont_inherit: Compiler options.
    optimize: Compiler options.

  Returns:
    The resulting code object.
  """
  # doctest passes some dummy string as the file name, AFAICT
  # but tf.function freaks-out if this doesn't look like a
  # python file name.
  del filename
  # Doctest always passes "single" here, you need exec for multiple lines.
  del mode

  source_ast = ast.parse(source)

  final = source_ast.body[-1]
  if isinstance(final, ast.Expr):
    # Wrap the final expression as `_print_if_not_none(expr)`
    print_it = ast.Expr(
        lineno=-1,
        col_offset=-1,
        value=ast.Call(
            func=ast.Name(
                id='_print_if_not_none',
                ctx=ast.Load(),
                lineno=-1,
                col_offset=-1),
            lineno=-1,
            col_offset=-1,
            args=[final],  # wrap the final Expression
            keywords=[]))
    source_ast.body[-1] = print_it

    # It's not clear why this step is necessary. `compile` is supposed to handle
    # AST directly.
    source = astor.to_source(source_ast)

  return compile(
      source,
      filename='dummy.py',
      mode='exec',
      flags=flags,
      dont_inherit=dont_inherit,
      optimize=optimize)