sdks/python/apache_beam/dataframe/pandas_docs_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.

"""A module for running the pandas docs (such as the users guide) against our
dataframe implementation.

Run as python -m apache_beam.dataframe.pandas_docs_test [getting_started ...]
"""

import argparse
import contextlib
import io
import multiprocessing
import os
import sys
import time
import urllib.request
import zipfile

from apache_beam.dataframe import doctests

PANDAS_VERSION = '1.1.1'
PANDAS_DIR = os.path.expanduser("~/.apache_beam/cache/pandas-" + PANDAS_VERSION)
PANDAS_DOCS_SOURCE = os.path.join(PANDAS_DIR, 'doc', 'source')


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '-p',
      '--parallel',
      type=int,
      default=0,
      help='Number of tests to run in parallel. '
      'Defaults to 0, meaning the number of cores on the machine.')
  parser.add_argument('docs', nargs='*')
  args = parser.parse_args()

  if not os.path.exists(PANDAS_DIR):
    # Download the pandas source.
    os.makedirs(os.path.dirname(PANDAS_DIR), exist_ok=True)
    zip = os.path.join(PANDAS_DIR + '.zip')
    if not os.path.exists(zip):
      url = (
          'https://github.com/pandas-dev/pandas/archive/v%s.zip' %
          PANDAS_VERSION)
      print('Downloading', url)
      with urllib.request.urlopen(url) as fin:
        with open(zip + '.tmp', 'wb') as fout:
          fout.write(fin.read())
        os.rename(zip + '.tmp', zip)

    print('Extracting', zip)
    with zipfile.ZipFile(zip, 'r') as handle:
      handle.extractall(os.path.dirname(PANDAS_DIR))

  tests = args.docs or ['getting_started', 'user_guide']
  paths = []
  filters = []

  # Explicit paths.
  for test in tests:
    if os.path.exists(test):
      paths.append(test)
    else:
      filters.append(test)

  # Names of pandas source files.
  for root, _, files in os.walk(PANDAS_DOCS_SOURCE):
    for name in files:
      if name.endswith('.rst'):
        path = os.path.join(root, name)
        if any(filter in path for filter in filters):
          paths.append(path)

  # Using a global here is a bit hacky, but avoids pickling issues when used
  # with multiprocessing.
  parallelism = max(args.parallel or multiprocessing.cpu_count(), len(paths))

  if parallelism > 1:
    pool_map = multiprocessing.pool.Pool(parallelism).imap_unordered
    run_tests = run_tests_capturing_stdout
    # Make sure slow tests get started first.
    paths.sort(
        key=lambda path: ('enhancingperf' in path, os.path.getsize(path)),
        reverse=True)
  else:
    pool_map = map
    run_tests = run_tests_streaming_stdout

  # Now run all the tests.
  running_summary = doctests.Summary()
  for count, (summary, stdout) in enumerate(pool_map(run_tests, paths)):
    running_summary += summary
    if stdout:
      print(stdout)
    print(count, '/', len(paths), 'done.')

  print('*' * 72)
  print("Final summary:")
  running_summary.summarize()


def run_tests_capturing_stdout(path):
  with deferred_stdout() as stdout:
    return run_tests(path), stdout()


def run_tests_streaming_stdout(path):
  return run_tests(path), None


def run_tests(path):
  # Optionally capture the stdout as interleaved test errors are painful
  # to debug.  On the other hand, if there is no parallelism, let the
  # output be streamed.
  start = time.time()
  with open(path) as f:
    rst = f.read()
  res = doctests.test_rst_ipython(
      rst,
      path,
      report=True,
      wont_implement_ok=['*'],
      not_implemented_ok=['*'],
      use_beam=False).summary
  print("Total time for {}: {:.2f} secs".format(path, time.time() - start))
  return res


@contextlib.contextmanager
def deferred_stdout():
  captured = io.StringIO()
  old_stdout, sys.stdout = sys.stdout, captured
  try:
    yield captured.getvalue
  finally:
    sys.stdout = old_stdout


if __name__ == '__main__':
  main()