# 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.

from typing import Dict, TYPE_CHECKING

import os
import re

import nbformat
import pytest

if TYPE_CHECKING:
    # pylint: disable=unused-import
    from typing import Any, List


def find_examples_jupyter_notebook_paths():
    examples_folder = os.path.dirname(__file__)
    for filename in os.listdir(examples_folder):
        if not filename.endswith('.ipynb'):
            continue
        yield os.path.join(examples_folder, filename)


@pytest.mark.parametrize('path', find_examples_jupyter_notebook_paths())
def test_can_run_examples_jupyter_notebook(path):
    assert_jupyter_notebook_has_working_code_cells(path)


def assert_jupyter_notebook_has_working_code_cells(path: str):
    """Checks that code cells in a Jupyter notebook actually run in sequence.

    State is kept between code cells. Imports and variables defined in one
    cell will be visible in later cells.
    """

    notebook = nbformat.read(path, nbformat.NO_CONVERT)
    state = {}  # type: Dict[str, Any]

    for cell in notebook.cells:
        if cell.cell_type == 'code':
            assert_code_cell_runs_and_prints_expected(cell, state)


def assert_code_cell_runs_and_prints_expected(cell: nbformat.NotebookNode,
                                              state: Dict):
    """Executes a code cell and compares captured output to saved output."""

    expected_outputs = cell.outputs[0].text.strip().split('\n')
    expected_lines = [canonicalize_printed_line(line)
                      for line in expected_outputs]

    output_lines = []  # type: List[str]

    def print_capture(*values, sep=' '):
        output_lines.extend(
                sep.join(str(e) for e in values).strip().split('\n'))

    state['print'] = print_capture
    exec(cell.source, state)

    actual_lines = [canonicalize_printed_line(line) for line in output_lines]

    assert len(actual_lines) == len(expected_lines)

    for i, line in enumerate(actual_lines):
        assert line == expected_lines[i]


def canonicalize_printed_line(line: str) -> str:
    """Replace highly precise numbers with less precise numbers.

    This allows the test to pass on machines that perform arithmetic slightly
    differently.

    Args:
        line: The line to canonicalize.

    Returns:
        The canonicalized line.
    """
    prev_end = 0
    result = []
    for match in re.finditer(r"[0-9]\.[0-9]+", line):
        start = match.start()
        end = match.end()
        result.append(line[prev_end:start])
        result.append(format(round(float(line[start:end]), 4), '.4f'))
        prev_end = end
    result.append(line[prev_end:])
    return ''.join(result).rstrip()


def test_canonicalize_printed_line():
    x = 'first 20.37378859061888 then 20.37378627319067'
    assert canonicalize_printed_line(x) == 'first 20.3738 then 20.3738'

@@ -1886,67 +1886,191 @@
     
-expected_outputs = cell.outputs%5B0%5D.text.strip().split('%5Cn')
+if cell.outputs and hasattr(cell.outputs%5B0%5D, 'text'):%0A        expected_outputs = cell.outputs%5B0%5D.text.strip().split('%5Cn')%0A    else:%0A        expected_outputs = %5B''%5D  # coverage: ignore
 %0A   
@@ -2334,24 +2334,16 @@
 values).
-strip().
 split('%5C
@@ -2392,16 +2392,39 @@
    exec(
+strip_magics_and_shows(
 cell.sou
@@ -2426,24 +2426,25 @@
 l.source
+)
 , state)
 %0A%0A    ac
@@ -2435,16 +2435,79 @@
 , state)
+%0A    output_lines = '%5Cn'.join(output_lines).strip().split('%5Cn')
 %0A%0A    ac
@@ -2719,16 +2719,411 @@
 es%5Bi%5D%0A%0A%0A
+def strip_magics_and_shows(text: str) -%3E str:%0A    %22%22%22Remove Jupyter magics and pyplot show commands.%22%22%22%0A    lines = %5Bline for line in text.split('%5Cn')%0A             if not contains_magic_or_show(line)%5D%0A    return '%5Cn'.join(lines)%0A%0A%0Adef contains_magic_or_show(line: str) -%3E bool:%0A    return (line.strip().startswith('%25') or%0A            'pyplot.show(' in line or%0A            'plt.show(' in line)%0A%0A%0A
 def cano
@@ -3502,16 +3502,17 @@
 0-9%5D
++
 %5C.%5B0-9%5D+
 %22, l
@@ -3507,16 +3507,27 @@
 %5C.%5B0-9%5D+
+(e-%5B0-9%5D+)?
 %22, line)
@@ -3962,8 +3962,120 @@
 0.3738'%0A
+    x = 'first 20.37e-8 then 224.373e-14'%0A    assert canonicalize_printed_line(x) == 'first 0.0000 then 0.0000'%0A

Make examples_test less brittle (#145)