<commit_before># Copyright 2013 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import logging
import os

from telemetry.core import exceptions
from telemetry.core import util
from telemetry import decorators
from telemetry.internal.backends.chrome import chrome_browser_backend
from telemetry.internal.backends.chrome import misc_web_contents_backend
from telemetry.internal import forwarders

class CrOSBrowserBackend(chrome_browser_backend.ChromeBrowserBackend):
  def __init__(self, cros_platform_backend, browser_options, cri, is_guest):
    super(CrOSBrowserBackend, self).__init__(
        cros_platform_backend, supports_tab_control=True,
        supports_extensions=not is_guest,
        browser_options=browser_options)
    assert browser_options.IsCrosBrowserOptions()
    # Initialize fields so that an explosion during init doesn't break in Close.
    self._cri = cri
    self._is_guest = is_guest
    self._forwarder = None
    self._remote_debugging_port = self._cri.GetRemotePort()
    self._port = self._remote_debugging_port

    extensions_to_load = browser_options.extensions_to_load

    # Copy extensions to temp directories on the device.
    # Note that we also perform this copy locally to ensure that
    # the owner of the extensions is set to chronos.
    for e in extensions_to_load:
      extension_dir = cri.RunCmdOnDevice(
          ['mktemp', '-d', '/tmp/extension_XXXXX'])[0].rstrip()
      e.local_path = os.path.join(extension_dir, os.path.basename(e.path))
      cri.PushFile(e.path, extension_dir)
      cri.Chown(extension_dir)

    self._cri.RestartUI(self.browser_options.clear_enterprise_policy)
    util.WaitFor(self.IsBrowserRunning, 20)

    # Delete test user's cryptohome vault (user data directory).
    if not self.browser_options.dont_override_profile:
      self._cri.RunCmdOnDevice(['cryptohome', '--action=remove', '--force',
                                '--user=%s' % self._username])

  @property
  def log_file_path(self):
    return None

  def GetBrowserStartupArgs(self):
    args = super(CrOSBrowserBackend, self).GetBrowserStartupArgs()
    args.extend([
            '--enable-smooth-scrolling',
            '--enable-threaded-compositing',
            # Allow devtools to connect to chrome.
            '--remote-debugging-port=%i' % self._remote_debugging_port,
            # Open a maximized window.
            '--start-maximized',
            # Disable system startup sound.
            '--ash-disable-system-sounds',
            # Ignore DMServer errors for policy fetches.
            '--allow-failed-policy-fetch-for-test',
            # Skip user image selection screen, and post login screens.
            '--oobe-skip-postlogin',
            # Debug logging.
            '--vmodule=*/chromeos/net/*=2,*/chromeos/login/*=2'])

    # Disable GAIA services unless we're using GAIA login, or if there's an
    # explicit request for it.
    if (self.browser_options.disable_gaia_services and
        not self.browser_options.gaia_login):
      args.append('--disable-gaia-services')

    trace_config_file = (self.platform_backend.tracing_controller_backend
                         .GetChromeTraceConfigFile())
    if trace_config_file:
      args.append('--trace-config-file=%s' % trace_config_file)

    return args

  @property
  def pid(self):
    return self._cri.GetChromePid()

  @property
  def browser_directory(self):
    result = self._cri.GetChromeProcess()
    if result and 'path' in result:
      return os.path.dirname(result['path'])
    return None

  @property
  def profile_directory(self):
    return '/home/chronos/Default'

  def __del__(self):
    self.Close()

  def Start(self):
    # Escape all commas in the startup arguments we pass to Chrome
    # because dbus-send delimits array elements by commas
    startup_args = [a.replace(',', '\\,') for a in self.GetBrowserStartupArgs()]

    # Restart Chrome with the login extension and remote debugging.
    logging.info('Restarting Chrome with flags and login')
    args = ['dbus-send', '--system', '--type=method_call',
            '--dest=org.chromium.SessionManager',
            '/org/chromium/SessionManager',
            'org.chromium.SessionManagerInterface.EnableChromeTesting',
            'boolean:true',
            'array:string:"%s"' % ','.join(startup_args)]
    logging.info('Starting Chrome %s', args)
    self._cri.RunCmdOnDevice(args)

    if not self._cri.local:
      # TODO(crbug.com/404771): Move port forwarding to network_controller.
      self._port = util.GetUnreservedAvailableLocalPort()
      self._forwarder = self._platform_backend.forwarder_factory.Create(
          forwarders.PortPairs(
              http=forwarders.PortPair(self._port, self._remote_debugging_port),
              https=None,
              dns=None), use_remote_port_forwarding=False)

    # Wait for oobe.
    self._WaitForBrowserToComeUp()
    self._InitDevtoolsClientBackend(
        remote_devtools_port=self._remote_debugging_port)
    util.WaitFor(lambda: self.oobe_exists, 30)

    if self.browser_options.auto_login:
      try:
        if self._is_guest:
          pid = self.pid
          self.oobe.NavigateGuestLogin()
          # Guest browsing shuts down the current browser and launches an
          # incognito browser in a separate process, which we need to wait for.
          util.WaitFor(lambda: pid != self.pid, 10)
        elif self.browser_options.gaia_login:
          self.oobe.NavigateGaiaLogin(self._username, self._password)
        else:
          self.oobe.NavigateFakeLogin(self._username, self._password,
              self._gaia_id, not self.browser_options.disable_gaia_services)

        self._WaitForLogin()
      except exceptions.TimeoutException:
        self._cri.TakeScreenshotWithPrefix('login-screen')
        raise exceptions.LoginException('Timed out going through login screen. '
                                        + self._GetLoginStatus())

    logging.info('Browser is up!')

  def Close(self):
    super(CrOSBrowserBackend, self).Close()

    if self._cri:
      self._cri.RestartUI(False) # Logs out.
      self._cri.CloseConnection()

    util.WaitFor(lambda: not self._IsCryptohomeMounted(), 180)

    if self._forwarder:
      self._forwarder.Close()
      self._forwarder = None

    if self._cri:
      for e in self._extensions_to_load:
        self._cri.RmRF(os.path.dirname(e.local_path))

    self._cri = None

  def IsBrowserRunning(self):
    return bool(self.pid)

  def GetStandardOutput(self):
    return 'Cannot get standard output on CrOS'

  def GetStackTrace(self):
    return (False, 'Cannot get stack trace on CrOS')

  def GetMostRecentMinidumpPath(self):
    return None

  def GetAllMinidumpPaths(self):
    return None

  def GetAllUnsymbolizedMinidumpPaths(self):
    return None

  def SymbolizeMinidump(self, minidump_path):
    return None

  @property
  @decorators.Cache
  def misc_web_contents_backend(self):
    """Access to chrome://oobe/login page."""
    return misc_web_contents_backend.MiscWebContentsBackend(self)

  @property
  def oobe(self):
    return self.misc_web_contents_backend.GetOobe()

  @property
  def oobe_exists(self):
    return self.misc_web_contents_backend.oobe_exists

  @property
  def _username(self):
    return self.browser_options.username

  @property
  def _password(self):
    return self.browser_options.password

  @property
  def _gaia_id(self):
    return self.browser_options.gaia_id

  def _IsCryptohomeMounted(self):
    username = '$guest' if self._is_guest else self._username
    return self._cri.IsCryptohomeMounted(username, self._is_guest)

  def _GetLoginStatus(self):
    """Returns login status. If logged in, empty string is returned."""
    status = ''
    if not self._IsCryptohomeMounted():
      status += 'Cryptohome not mounted. '
    if not self.HasBrowserFinishedLaunching():
      status += 'Browser didn\'t launch. '
    if self.oobe_exists:
      status += 'OOBE not dismissed.'
    return status

  def _IsLoggedIn(self):
    """Returns True if cryptohome has mounted, the browser is
    responsive to devtools requests, and the oobe has been dismissed."""
    return not self._GetLoginStatus()

  def _WaitForLogin(self):
    # Wait for cryptohome to mount.
    util.WaitFor(self._IsLoggedIn, 60)

    # For incognito mode, the session manager actually relaunches chrome with
    # new arguments, so we have to wait for the browser to come up.
    self._WaitForBrowserToComeUp()

    # Wait for extensions to load.
    if self._supports_extensions:
      self._WaitForExtensionsToLoad()
<commit_msg>Speculative fix for a possible race in browser restart.<commit_after># Copyright 2013 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import logging
import os

from telemetry.core import exceptions
from telemetry.core import util
from telemetry import decorators
from telemetry.internal.backends.chrome import chrome_browser_backend
from telemetry.internal.backends.chrome import misc_web_contents_backend
from telemetry.internal import forwarders

class CrOSBrowserBackend(chrome_browser_backend.ChromeBrowserBackend):
  def __init__(self, cros_platform_backend, browser_options, cri, is_guest):
    super(CrOSBrowserBackend, self).__init__(
        cros_platform_backend, supports_tab_control=True,
        supports_extensions=not is_guest,
        browser_options=browser_options)
    assert browser_options.IsCrosBrowserOptions()
    # Initialize fields so that an explosion during init doesn't break in Close.
    self._cri = cri
    self._is_guest = is_guest
    self._forwarder = None
    self._remote_debugging_port = self._cri.GetRemotePort()
    self._port = self._remote_debugging_port

    extensions_to_load = browser_options.extensions_to_load

    # Copy extensions to temp directories on the device.
    # Note that we also perform this copy locally to ensure that
    # the owner of the extensions is set to chronos.
    for e in extensions_to_load:
      extension_dir = cri.RunCmdOnDevice(
          ['mktemp', '-d', '/tmp/extension_XXXXX'])[0].rstrip()
      e.local_path = os.path.join(extension_dir, os.path.basename(e.path))
      cri.PushFile(e.path, extension_dir)
      cri.Chown(extension_dir)

    self._cri.RestartUI(self.browser_options.clear_enterprise_policy)
    util.WaitFor(self.IsBrowserRunning, 20)

    # Delete test user's cryptohome vault (user data directory).
    if not self.browser_options.dont_override_profile:
      self._cri.RunCmdOnDevice(['cryptohome', '--action=remove', '--force',
                                '--user=%s' % self._username])

  @property
  def log_file_path(self):
    return None

  def GetBrowserStartupArgs(self):
    args = super(CrOSBrowserBackend, self).GetBrowserStartupArgs()

    logging_patterns = ['*/chromeos/net/*',
                        '*/chromeos/login/*',
                        'application_lifetime',
                        'chrome_browser_main_posix']
    vmodule = '--vmodule='
    for pattern in logging_patterns:
      vmodule += '%s=2,' % pattern
    vmodule = vmodule.rstrip(',')

    args.extend([
            '--enable-smooth-scrolling',
            '--enable-threaded-compositing',
            # Allow devtools to connect to chrome.
            '--remote-debugging-port=%i' % self._remote_debugging_port,
            # Open a maximized window.
            '--start-maximized',
            # Disable system startup sound.
            '--ash-disable-system-sounds',
            # Ignore DMServer errors for policy fetches.
            '--allow-failed-policy-fetch-for-test',
            # Skip user image selection screen, and post login screens.
            '--oobe-skip-postlogin',
            # Debug logging.
            vmodule])

    # Disable GAIA services unless we're using GAIA login, or if there's an
    # explicit request for it.
    if (self.browser_options.disable_gaia_services and
        not self.browser_options.gaia_login):
      args.append('--disable-gaia-services')

    trace_config_file = (self.platform_backend.tracing_controller_backend
                         .GetChromeTraceConfigFile())
    if trace_config_file:
      args.append('--trace-config-file=%s' % trace_config_file)

    return args

  @property
  def pid(self):
    return self._cri.GetChromePid()

  @property
  def browser_directory(self):
    result = self._cri.GetChromeProcess()
    if result and 'path' in result:
      return os.path.dirname(result['path'])
    return None

  @property
  def profile_directory(self):
    return '/home/chronos/Default'

  def __del__(self):
    self.Close()

  def Start(self):
    # Escape all commas in the startup arguments we pass to Chrome
    # because dbus-send delimits array elements by commas
    startup_args = [a.replace(',', '\\,') for a in self.GetBrowserStartupArgs()]

    # Restart Chrome with the login extension and remote debugging.
    pid = self.pid
    logging.info('Restarting Chrome (pid=%d) with remote port', pid)
    args = ['dbus-send', '--system', '--type=method_call',
            '--dest=org.chromium.SessionManager',
            '/org/chromium/SessionManager',
            'org.chromium.SessionManagerInterface.EnableChromeTesting',
            'boolean:true',
            'array:string:"%s"' % ','.join(startup_args)]
    logging.info('Starting Chrome %s', args)
    self._cri.RunCmdOnDevice(args)

    if not self._cri.local:
      # TODO(crbug.com/404771): Move port forwarding to network_controller.
      self._port = util.GetUnreservedAvailableLocalPort()
      self._forwarder = self._platform_backend.forwarder_factory.Create(
          forwarders.PortPairs(
              http=forwarders.PortPair(self._port, self._remote_debugging_port),
              https=None,
              dns=None), use_remote_port_forwarding=False)

    # Wait for new chrome and oobe.
    util.WaitFor(lambda: pid != self.pid, 15)
    self._WaitForBrowserToComeUp()
    self._InitDevtoolsClientBackend(
        remote_devtools_port=self._remote_debugging_port)
    util.WaitFor(lambda: self.oobe_exists, 30)

    if self.browser_options.auto_login:
      try:
        if self._is_guest:
          pid = self.pid
          self.oobe.NavigateGuestLogin()
          # Guest browsing shuts down the current browser and launches an
          # incognito browser in a separate process, which we need to wait for.
          util.WaitFor(lambda: pid != self.pid, 15)
        elif self.browser_options.gaia_login:
          self.oobe.NavigateGaiaLogin(self._username, self._password)
        else:
          self.oobe.NavigateFakeLogin(self._username, self._password,
              self._gaia_id, not self.browser_options.disable_gaia_services)

        self._WaitForLogin()
      except exceptions.TimeoutException:
        self._cri.TakeScreenshotWithPrefix('login-screen')
        raise exceptions.LoginException('Timed out going through login screen. '
                                        + self._GetLoginStatus())

    logging.info('Browser is up!')

  def Close(self):
    super(CrOSBrowserBackend, self).Close()

    if self._cri:
      self._cri.RestartUI(False) # Logs out.
      self._cri.CloseConnection()

    util.WaitFor(lambda: not self._IsCryptohomeMounted(), 180)

    if self._forwarder:
      self._forwarder.Close()
      self._forwarder = None

    if self._cri:
      for e in self._extensions_to_load:
        self._cri.RmRF(os.path.dirname(e.local_path))

    self._cri = None

  def IsBrowserRunning(self):
    return bool(self.pid)

  def GetStandardOutput(self):
    return 'Cannot get standard output on CrOS'

  def GetStackTrace(self):
    return (False, 'Cannot get stack trace on CrOS')

  def GetMostRecentMinidumpPath(self):
    return None

  def GetAllMinidumpPaths(self):
    return None

  def GetAllUnsymbolizedMinidumpPaths(self):
    return None

  def SymbolizeMinidump(self, minidump_path):
    return None

  @property
  @decorators.Cache
  def misc_web_contents_backend(self):
    """Access to chrome://oobe/login page."""
    return misc_web_contents_backend.MiscWebContentsBackend(self)

  @property
  def oobe(self):
    return self.misc_web_contents_backend.GetOobe()

  @property
  def oobe_exists(self):
    return self.misc_web_contents_backend.oobe_exists

  @property
  def _username(self):
    return self.browser_options.username

  @property
  def _password(self):
    return self.browser_options.password

  @property
  def _gaia_id(self):
    return self.browser_options.gaia_id

  def _IsCryptohomeMounted(self):
    username = '$guest' if self._is_guest else self._username
    return self._cri.IsCryptohomeMounted(username, self._is_guest)

  def _GetLoginStatus(self):
    """Returns login status. If logged in, empty string is returned."""
    status = ''
    if not self._IsCryptohomeMounted():
      status += 'Cryptohome not mounted. '
    if not self.HasBrowserFinishedLaunching():
      status += 'Browser didn\'t launch. '
    if self.oobe_exists:
      status += 'OOBE not dismissed.'
    return status

  def _IsLoggedIn(self):
    """Returns True if cryptohome has mounted, the browser is
    responsive to devtools requests, and the oobe has been dismissed."""
    return not self._GetLoginStatus()

  def _WaitForLogin(self):
    # Wait for cryptohome to mount.
    util.WaitFor(self._IsLoggedIn, 60)

    # For incognito mode, the session manager actually relaunches chrome with
    # new arguments, so we have to wait for the browser to come up.
    self._WaitForBrowserToComeUp()

    # Wait for extensions to load.
    if self._supports_extensions:
      self._WaitForExtensionsToLoad()
<|endoftext|>


<commit_before>#    Copyright 2015 ARM Limited
#
# 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.
#
# pylint: disable=no-name-in-module,import-error,no-member

import os
import subprocess
from distutils.version import StrictVersion

# pylint: disable=wrong-import-position
import_error_str = ''
try:
    import IPython
except ImportError as import_error:
    IPython = None
    # Importing IPython can fail for a variety of reasons, report the actual
    # failure unless it's just that the package is not present
    if import_error.message.startswith("No module named"):  # pylint: disable=E1101
        import_error_str = 'ipynb_exporter requires ipython package to be installed'
    else:
        import_error_str = import_error.message  # pylint: disable=redefined-variable-type

# The current code generates notebooks version 3
NBFORMAT_VERSION = 3

if IPython:
    if StrictVersion('4.0.0') > StrictVersion(IPython.__version__) >= StrictVersion('3.0.0'):
        import IPython.kernel
        import IPython.nbformat

        def read_notebook(notebook_in):
            return IPython.nbformat.reads(notebook_in, NBFORMAT_VERSION)  # pylint: disable=E1101

        def write_notebook(notebook, fout):
            IPython.nbformat.write(notebook, fout)  # pylint: disable=E1101

        NotebookNode = IPython.nbformat.NotebookNode  # pylint: disable=E1101

        IPYTHON_NBCONVERT_HTML = ['ipython', 'nbconvert', '--to=html']
        IPYTHON_NBCONVERT_PDF = ['ipython', 'nbconvert', '--to=pdf']
    elif StrictVersion('3.0.0') > StrictVersion(IPython.__version__) >= StrictVersion('2.0.0'):
        import IPython.kernel
        import IPython.nbformat.v3

        def read_notebook(notebook_in):
            return IPython.nbformat.v3.reads_json(notebook_in)  # pylint: disable=E1101

        def write_notebook(notebook, fout):
            IPython.nbformat.v3.nbjson.JSONWriter().write(notebook, fout)  # pylint: disable=E1101

        NotebookNode = IPython.nbformat.v3.NotebookNode  # pylint: disable=E1101

        IPYTHON_NBCONVERT_HTML = ['ipython', 'nbconvert', '--to=html']
        IPYTHON_NBCONVERT_PDF = ['ipython', 'nbconvert', '--to=latex',
                                 '--post=PDF']
    else:
        # Unsupported IPython version
        import_error_str = 'Unsupported IPython version {}'.format(IPython.__version__)

def parse_valid_output(msg):
    """Parse a valid result from an execution of a cell in an ipython kernel"""
    msg_type = msg["msg_type"]
    if msg_type == 'error':
        msg_type = 'pyerr'
    elif msg_type == 'execute_result':
        msg_type = 'pyout'

    content = msg["content"]
    out = NotebookNode(output_type=msg_type)

    if msg_type == "stream":
        out.stream = content["name"]
        try:
            out.text = content['data']
        except KeyError:
            out.text = content['text']
    elif msg_type in ("display_data", "pyout"):
        for mime, data in content["data"].iteritems():
            if mime == "text/plain":
                attr = "text"
            else:
                attr = mime.split("/")[-1]
            setattr(out, attr, data)
    elif msg_type == "pyerr":
        out.ename = content["ename"]
        out.evalue = content["evalue"]
        out.traceback = content["traceback"]
    else:
        raise ValueError("Unknown msg_type {}".format(msg_type))

    return out

def run_cell(kernel_client, cell):
    """Run a cell of a notebook in an ipython kernel and return its output"""
    kernel_client.execute(cell.input)

    input_acknowledged = False
    outs = []
    while True:
        msg = kernel_client.get_iopub_msg()

        if msg["msg_type"] == "status":
            if msg["content"]["execution_state"] == "idle" and input_acknowledged:
                break
        elif msg["msg_type"] in ('pyin', 'execute_input'):
            input_acknowledged = True
        else:
            out = parse_valid_output(msg)
            outs.append(out)

    return outs

def run_notebook(notebook):
    """Run the notebook"""

    kernel_manager = IPython.kernel.KernelManager()
    kernel_manager.start_kernel(stderr=open(os.devnull, 'w'))
    kernel_client = kernel_manager.client()
    kernel_client.start_channels()

    for sheet in notebook.worksheets:
        for (prompt_number, cell) in enumerate(sheet.cells, 1):
            if cell.cell_type != "code":
                continue

            cell.outputs = run_cell(kernel_client, cell)

            cell.prompt_number = prompt_number
            if cell.outputs and cell.outputs[0]['output_type'] == 'pyout':
                cell.outputs[0]["prompt_number"] = prompt_number

    kernel_manager.shutdown_kernel()

def export_notebook(nbbasename, output_directory, output_format):
    """Generate a PDF or HTML from the ipython notebook

    output_format has to be either 'pdf' or 'html'.  These are the
    only formats currently supported.

    ipython nbconvert claims that the CLI is not stable, so keep this
    function here to be able to cope with inconsistencies

    """

    if output_format == "html":
        ipython_command = IPYTHON_NBCONVERT_HTML
    elif output_format == "pdf":
        ipython_command = IPYTHON_NBCONVERT_PDF
    else:
        raise ValueError("Unknown output format: {}".format(output_format))

    prev_dir = os.getcwd()
    os.chdir(output_directory)

    with open(os.devnull, 'w') as devnull:
        subprocess.check_call(ipython_command + [nbbasename], stderr=devnull)

    os.chdir(prev_dir)
<commit_msg>ipython: Updated to work with the latest ipython version<commit_after>#    Copyright 2015 ARM Limited
#
# 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.
#
# pylint: disable=no-name-in-module,import-error,no-member

import os
import subprocess
from distutils.version import StrictVersion

# pylint: disable=wrong-import-position,ungrouped-imports
import_error_str = ''
try:
    import IPython
except ImportError as import_error:
    IPython = None
    # Importing IPython can fail for a variety of reasons, report the actual
    # failure unless it's just that the package is not present
    if import_error.message.startswith("No module named"):  # pylint: disable=E1101
        import_error_str = 'ipynb_exporter requires ipython package to be installed'
    else:
        import_error_str = import_error.message  # pylint: disable=redefined-variable-type

# The current code generates notebooks version 3
NBFORMAT_VERSION = 3

if IPython:
    if StrictVersion('5.0.0') > StrictVersion(IPython.__version__) >= StrictVersion('4.0.0'):
        import nbformat
        from jupyter_client.manager import KernelManager

        def read_notebook(notebook_in):  # pylint: disable=function-redefined
            return nbformat.reads(notebook_in, NBFORMAT_VERSION)  # pylint: disable=E1101

        def write_notebook(notebook, fout):  # pylint: disable=function-redefined
            nbformat.write(notebook, fout)  # pylint: disable=E1101

        NotebookNode = nbformat.NotebookNode  # pylint: disable=E1101

        IPYTHON_NBCONVERT_HTML = ['jupyter', 'nbconvert', '--to html']
        IPYTHON_NBCONVERT_PDF = ['jupyter', 'nbconvert', '--to pdf']

    elif StrictVersion('4.0.0') > StrictVersion(IPython.__version__) >= StrictVersion('3.0.0'):
        from IPython.kernel import KernelManager
        import IPython.nbformat

        def read_notebook(notebook_in):  # pylint: disable=function-redefined
            return IPython.nbformat.reads(notebook_in, NBFORMAT_VERSION)  # pylint: disable=E1101

        def write_notebook(notebook, fout):  # pylint: disable=function-redefined
            IPython.nbformat.write(notebook, fout)  # pylint: disable=E1101

        NotebookNode = IPython.nbformat.NotebookNode  # pylint: disable=E1101

        IPYTHON_NBCONVERT_HTML = ['ipython', 'nbconvert', '--to=html']
        IPYTHON_NBCONVERT_PDF = ['ipython', 'nbconvert', '--to=pdf']
    elif StrictVersion('3.0.0') > StrictVersion(IPython.__version__) >= StrictVersion('2.0.0'):
        from IPython.kernel import KernelManager
        import IPython.nbformat.v3

        def read_notebook(notebook_in):  # pylint: disable=function-redefined
            return IPython.nbformat.v3.reads_json(notebook_in)  # pylint: disable=E1101

        def write_notebook(notebook, fout):  # pylint: disable=function-redefined
            IPython.nbformat.v3.nbjson.JSONWriter().write(notebook, fout)  # pylint: disable=E1101

        NotebookNode = IPython.nbformat.v3.NotebookNode  # pylint: disable=E1101

        IPYTHON_NBCONVERT_HTML = ['ipython', 'nbconvert', '--to=html']
        IPYTHON_NBCONVERT_PDF = ['ipython', 'nbconvert', '--to=latex',
                                 '--post=PDF']
    else:
        # Unsupported IPython version
        import_error_str = 'Unsupported IPython version {}'.format(IPython.__version__)

def parse_valid_output(msg):
    """Parse a valid result from an execution of a cell in an ipython kernel"""
    msg_type = msg["msg_type"]
    if msg_type == 'error':
        msg_type = 'pyerr'
    elif msg_type == 'execute_result':
        msg_type = 'pyout'

    content = msg["content"]
    out = NotebookNode(output_type=msg_type)

    if msg_type == "stream":
        out.stream = content["name"]
        try:
            out.text = content['data']
        except KeyError:
            out.text = content['text']
    elif msg_type in ("display_data", "pyout"):
        for mime, data in content["data"].iteritems():
            if mime == "text/plain":
                attr = "text"
            else:
                attr = mime.split("/")[-1]
            setattr(out, attr, data)
    elif msg_type == "pyerr":
        out.ename = content["ename"]
        out.evalue = content["evalue"]
        out.traceback = content["traceback"]
    else:
        raise ValueError("Unknown msg_type {}".format(msg_type))

    return out

def run_cell(kernel_client, cell):
    """Run a cell of a notebook in an ipython kernel and return its output"""
    kernel_client.execute(cell.input)

    input_acknowledged = False
    outs = []
    while True:
        msg = kernel_client.get_iopub_msg()

        if msg["msg_type"] == "status":
            if msg["content"]["execution_state"] == "idle" and input_acknowledged:
                break
        elif msg["msg_type"] in ('pyin', 'execute_input'):
            input_acknowledged = True
        else:
            out = parse_valid_output(msg)
            outs.append(out)

    return outs

def run_notebook(notebook):
    """Run the notebook"""
    kernel_manager = KernelManager()
    kernel_manager.start_kernel(stderr=open(os.devnull, 'w'))
    kernel_client = kernel_manager.client()
    kernel_client.start_channels()

    for sheet in notebook.worksheets:
        for (prompt_number, cell) in enumerate(sheet.cells, 1):
            if cell.cell_type != "code":
                continue

            cell.outputs = run_cell(kernel_client, cell)

            cell.prompt_number = prompt_number
            if cell.outputs and cell.outputs[0]['output_type'] == 'pyout':
                cell.outputs[0]["prompt_number"] = prompt_number

    kernel_manager.shutdown_kernel()

def export_notebook(nbbasename, output_directory, output_format):
    """Generate a PDF or HTML from the ipython notebook

    output_format has to be either 'pdf' or 'html'.  These are the
    only formats currently supported.

    ipython nbconvert claims that the CLI is not stable, so keep this
    function here to be able to cope with inconsistencies

    """

    if output_format == "html":
        ipython_command = IPYTHON_NBCONVERT_HTML
    elif output_format == "pdf":
        ipython_command = IPYTHON_NBCONVERT_PDF
    else:
        raise ValueError("Unknown output format: {}".format(output_format))

    prev_dir = os.getcwd()
    os.chdir(output_directory)

    with open(os.devnull, 'w') as devnull:
        subprocess.check_call(ipython_command + [nbbasename], stderr=devnull)

    os.chdir(prev_dir)
<|endoftext|>


<commit_before># Copyright (c) 2014 by Ecreall under licence AGPL terms
# avalaible on http://www.gnu.org/licenses/agpl.html

# licence: AGPL
# author: Amen Souissi

import time
import transaction
import datetime

from pyramid.threadlocal import get_current_registry

import zmq
from zmq.eventloop.zmqstream import ZMQStream

from .core import FlowNode, BehavioralFlowNode, ProcessFinished
from dace.util import (
    get_system_request,
    get_socket,
    DelayedCallback,
    get_zmq_context,
    EventJob)

callbacks = {}

def push_callback_after_commit(event, callback, callback_params, deadline):
    # Create job object now before the end of the interaction so we have
    # the logged in user.
    job = EventJob('system')
    def after_commit_hook(status, *args, **kws):
        # status is true if the commit succeeded, or false if the commit aborted.
        if status:
            # Set callable now, we are sure to have p_oid
            event, callback, callback_params, deadline, job = args
            job.callable = callback
            job.args = callback_params
            dc = DelayedCallback(job, deadline, event._p_oid)
            dc.start()

    transaction.get().addAfterCommitHook(after_commit_hook,
            args=(event, callback, callback_params, deadline, job))

class Event(BehavioralFlowNode, FlowNode):

    def __init__(self, definition, eventKind, **kwargs):
        super(Event, self).__init__(definition, **kwargs)
        self.eventKind = eventKind
        self.execution_prepared = False
        self.execution_finished = False
        if eventKind:
            eventKind.event = self

    def __call__(self, transition):
        pass

    def prepare_for_execution(self):
        self.execution_prepared = True
        self.execution_finished = False

    def start(self, transaction):
        wi = self._get_workitem()
        if wi is not None:
            super(Event, self).start(transaction)
            self.execute()
            self.finish_behavior(wi)
            self.execution_finished = True
            self.execution_prepared = False
        else:
            self.stop()

    def replay_path(self, decision, transaction):
        if not self.execution_finished:
            self.execute()
            self.finish_behavior(decision)
            self.execution_finished = True
            self.execution_prepared = False

        super(Event, self).replay_path(decision, transaction)

    def validate(self):
        pass # pragma: no cover

    def execute(self):
        pass # pragma: no cover

    def stop(self):
        self.execution_prepared = False

class Throwing(Event):

    def validate(self):
        return True

    def prepare_for_execution(self):
        if not self.execution_prepared:
            super(Throwing, self).prepare_for_execution()
            if self.validate():
                self.start(None)
            else:
                self.stop()

    def execute(self):
        if self.eventKind is None:
            return

        return self.eventKind.execute()

class Catching(Event):

    def validate(self):
        if self.eventKind is None:
            return True

        return self.eventKind.validate()

    def prepare_for_execution(self):
        # If it's a empty StartEvent, execute the callback directly.
        if not self.execution_prepared:
            super(Catching, self).prepare_for_execution()
            if self.eventKind is None:
                if self.validate():
                    self.start(None)
                else:
                    self.stop()
            else:
                self.eventKind.prepare_for_execution()

    def stop(self):
        super(Catching, self).stop()
        if self.eventKind is not None:
            self.eventKind.stop()

class StartEvent(Catching):

    def start(self, transaction):
        if not self.process._started:
            wi = self._get_workitem()
            if wi is not None:
                super(StartEvent, self).start(transaction)
                self.process.start()
                #self.finish_behavior(wi)
                self.execution_finished = True
            else:
                self.stop()
        else:
            self.stop()

class IntermediateThrowEvent(Throwing):
    pass # pragma: no cover

class IntermediateCatchEvent(Catching):
    pass # pragma: no cover

class EndEvent(Throwing):

    def finish_behavior(self, work_item):
        super(EndEvent, self).finish_behavior(work_item)
        if isinstance(self.eventKind, TerminateEvent):
            return
        # Remove all workitems from process
        for node in self.process.nodes:
            try:
                node.stop()
                node.setproperty('workitems', [])
            except:
                pass

        self.process._finished = True
        registry = get_current_registry()
        registry.notify(ProcessFinished(self))
        current_process = self.process
        if current_process.definition.isSubProcess:
            request = get_system_request()
            root_process = current_process.attachedTo.process
            current_process.attachedTo.finish_execution(None, request)
            root_process.execution_context.remove_sub_execution_context(
                                          current_process.execution_context)
            parent_node = current_process.attachedTo.node
            if current_process in parent_node.sub_processes:
                parent_node.sub_processes.remove(current_process)

        if current_process.definition.isVolatile and \
           getattr(current_process, '__property__', None):
            getattr(current_process.__parent__.__class__,
                    current_process.__property__).remove(current_process.__parent__,
                                                      current_process)

class EventKind(object):

    # cette operation est appelee par les evenements "Catching"
    def validate(self):
        return True

    def prepare_for_execution(self, restart=False):
        pass # pragma: no cover

    # cette operation est appelee par les evenements "Throwing"
    def execute(self):
        pass # pragma: no cover

    def stop(self):
        pass # pragma: no cover

    @property
    def definition(self):
        return self.event.definition.eventKind

class ConditionalEvent(EventKind):

    def validate(self):
        return self.definition.condition(self.event.process)

    def prepare_for_execution(self, restart=False):
        self._push_callback()

    def _callback(self):
        if self.event._p_oid in callbacks:
            del callbacks[self.event._p_oid]

        if self.validate():
            #log.info('%s %s', self.event, "validate ok")
            wi = self.event._get_workitem()
            if wi is not None:
                self.event.start(None)
            else:
                self.stop()
        else:
            #log.info('%s %s', self.event, "validate ko, retry in 1s")
            self._push_callback()

    def _push_callback(self):
        push_callback_after_commit(self.event, self._callback, (), 1000)

    def stop(self):
        get_socket().send_pyobj(('stop', self.event._p_oid))

class TerminateEvent(EventKind):
    pass

# https://github.com/zeromq/pyzmq/blob/master/examples/poll/pubsub.py
# http://learning-0mq-with-pyzmq.readthedocs.org/en/latest/pyzmq/multisocket/tornadoeventloop.html

def get_signal_socket_url():
    return 'tcp://127.0.0.1:12346'

class Listener(object):
    def __init__(self, job, identifier):
        self.job = job
        self.identifier = identifier

    def start(self):
        identifier = self.identifier
        job = self.job
        def execute_next(msg):
            # We can't use zodb object from outside here because
            # this code is executed in another thread (eventloop)
            # We don't have site or interaction, so the job must be created
            # before.
            # we can't use push_callback_after_commit here because
            # it will never commit in this thread (eventloop)
            if identifier in callbacks:
                callbacks[identifier].close()
                del callbacks[identifier]

            job.args = (msg, )
            # wait 2s that the throw event transaction has committed
            dc = DelayedCallback(job, 2000)
            dc.start()

        ctx = get_zmq_context()
        s = ctx.socket(zmq.SUB)
        s.setsockopt_string(zmq.SUBSCRIBE, u'')
        s.connect(get_signal_socket_url())
        stream = ZMQStream(s)
        callbacks[identifier] = stream
        stream.on_recv(execute_next)

class SignalEvent(EventKind):

    _msg = None

    def validate(self):
        return self.definition.refSignal(self.event.process) == self._msg

    def prepare_for_execution(self, restart=False):
        job = EventJob('system')
        transaction.commit()  # needed to have self.event._p_oid
        job.callable = self._callback
        event_oid = self.event._p_oid
        listener = Listener(job, event_oid)
        get_socket().send_pyobj(('start', listener))

    def stop(self):
        # close ZMQStream
        get_socket().send_pyobj(('close', self.event._p_oid))

    def _callback(self, msg):
        # convert str (py27) / bytes (py34) to unicode (py27) or str (py34)
        self._msg = msg[0].decode('utf-8')
        if self.validate():
            wi = self.event._get_workitem()
            if wi is not None:
                self.event.start(None)
            else:
                self.event.stop()

    def execute(self):
        ref = self.definition.refSignal(self.event.process)
        ctx = get_zmq_context()
        s = ctx.socket(zmq.PUB)
        s.bind(get_signal_socket_url())
        # Sleep to allow sockets to connect.
        time.sleep(0.2)
        s.send_string(ref)
        s.close()

class TimerEvent(EventKind):

    def __init__(self):
        super(TimerEvent, self).__init__()
        self.time_date = None
        self.time_duration = None
        self.time_cycle = None

    def _prepare_time(self, time, restart=False):
        if getattr(self, time, None) is None or not restart:
            setattr(self, time, getattr(self.definition, time)(self.event.process))#TODO

    def _start_time(self, restart=False):
        """Return start time in milliseconds.
        """
        if self.definition.time_date is not None:
            self._prepare_time('time_date', restart)
            if not isinstance(self.time_date, datetime.datetime):
                raise TypeError("Unsupported deadline %r" % self.time_date)
            return (time.mktime(self.time_date.timetuple()) - time.time()) * 1000
        elif self.definition.time_duration is not None:
            self._prepare_time('time_duration', restart)
            if not isinstance(self.time_duration, datetime.timedelta):
                raise TypeError("Unsupported deadline %r" % self.time_duration)
            return self.time_duration.total_seconds() * 1000
        elif self.definition.time_cycle is not None:
            raise NotImplementedError
        else:
            raise TypeError("Unsupported deadline")

    def validate(self):
        return True

    def prepare_for_execution(self, restart=False):
        self._push_callback(restart)

    def _callback(self):
        if self.event._p_oid in callbacks:
            del callbacks[self.event._p_oid]

        if self.validate():
            wi = self.event._get_workitem()
            if wi is not None:
                self.event.start(None)
            else:
                self.event.stop()
        else:
            self._push_callback()

    def _push_callback(self, restart=False):
        deadline = self._start_time(restart)
        push_callback_after_commit(self.event, self._callback, (), deadline)

    def stop(self):
        # stop DelayedCallback
        get_socket().send_pyobj(('stop', self.event._p_oid))
<commit_msg>rename push_callback_after_commit to push_event_callback_after_commit to avoid ambiguity<commit_after># Copyright (c) 2014 by Ecreall under licence AGPL terms
# avalaible on http://www.gnu.org/licenses/agpl.html

# licence: AGPL
# author: Amen Souissi

import time
import transaction
import datetime

from pyramid.threadlocal import get_current_registry

import zmq
from zmq.eventloop.zmqstream import ZMQStream

from .core import FlowNode, BehavioralFlowNode, ProcessFinished
from dace.util import (
    get_system_request,
    get_socket,
    DelayedCallback,
    get_zmq_context,
    EventJob)

callbacks = {}

def push__event_callback_after_commit(event, callback, callback_params, deadline):
    # Create job object now before the end of the interaction so we have
    # the logged in user.
    job = EventJob('system')
    def after_commit_hook(status, *args, **kws):
        # status is true if the commit succeeded, or false if the commit aborted.
        if status:
            # Set callable now, we are sure to have p_oid
            event, callback, callback_params, deadline, job = args
            job.callable = callback
            job.args = callback_params
            dc = DelayedCallback(job, deadline, event._p_oid)
            dc.start()

    transaction.get().addAfterCommitHook(after_commit_hook,
            args=(event, callback, callback_params, deadline, job))

class Event(BehavioralFlowNode, FlowNode):

    def __init__(self, definition, eventKind, **kwargs):
        super(Event, self).__init__(definition, **kwargs)
        self.eventKind = eventKind
        self.execution_prepared = False
        self.execution_finished = False
        if eventKind:
            eventKind.event = self

    def __call__(self, transition):
        pass

    def prepare_for_execution(self):
        self.execution_prepared = True
        self.execution_finished = False

    def start(self, transaction):
        wi = self._get_workitem()
        if wi is not None:
            super(Event, self).start(transaction)
            self.execute()
            self.finish_behavior(wi)
            self.execution_finished = True
            self.execution_prepared = False
        else:
            self.stop()

    def replay_path(self, decision, transaction):
        if not self.execution_finished:
            self.execute()
            self.finish_behavior(decision)
            self.execution_finished = True
            self.execution_prepared = False

        super(Event, self).replay_path(decision, transaction)

    def validate(self):
        pass # pragma: no cover

    def execute(self):
        pass # pragma: no cover

    def stop(self):
        self.execution_prepared = False

class Throwing(Event):

    def validate(self):
        return True

    def prepare_for_execution(self):
        if not self.execution_prepared:
            super(Throwing, self).prepare_for_execution()
            if self.validate():
                self.start(None)
            else:
                self.stop()

    def execute(self):
        if self.eventKind is None:
            return

        return self.eventKind.execute()

class Catching(Event):

    def validate(self):
        if self.eventKind is None:
            return True

        return self.eventKind.validate()

    def prepare_for_execution(self):
        # If it's a empty StartEvent, execute the callback directly.
        if not self.execution_prepared:
            super(Catching, self).prepare_for_execution()
            if self.eventKind is None:
                if self.validate():
                    self.start(None)
                else:
                    self.stop()
            else:
                self.eventKind.prepare_for_execution()

    def stop(self):
        super(Catching, self).stop()
        if self.eventKind is not None:
            self.eventKind.stop()

class StartEvent(Catching):

    def start(self, transaction):
        if not self.process._started:
            wi = self._get_workitem()
            if wi is not None:
                super(StartEvent, self).start(transaction)
                self.process.start()
                #self.finish_behavior(wi)
                self.execution_finished = True
            else:
                self.stop()
        else:
            self.stop()

class IntermediateThrowEvent(Throwing):
    pass # pragma: no cover

class IntermediateCatchEvent(Catching):
    pass # pragma: no cover

class EndEvent(Throwing):

    def finish_behavior(self, work_item):
        super(EndEvent, self).finish_behavior(work_item)
        if isinstance(self.eventKind, TerminateEvent):
            return
        # Remove all workitems from process
        for node in self.process.nodes:
            try:
                node.stop()
                node.setproperty('workitems', [])
            except:
                pass

        self.process._finished = True
        registry = get_current_registry()
        registry.notify(ProcessFinished(self))
        current_process = self.process
        if current_process.definition.isSubProcess:
            request = get_system_request()
            root_process = current_process.attachedTo.process
            current_process.attachedTo.finish_execution(None, request)
            root_process.execution_context.remove_sub_execution_context(
                                          current_process.execution_context)
            parent_node = current_process.attachedTo.node
            if current_process in parent_node.sub_processes:
                parent_node.sub_processes.remove(current_process)

        if current_process.definition.isVolatile and \
           getattr(current_process, '__property__', None):
            getattr(current_process.__parent__.__class__,
                    current_process.__property__).remove(current_process.__parent__,
                                                      current_process)

class EventKind(object):

    # cette operation est appelee par les evenements "Catching"
    def validate(self):
        return True

    def prepare_for_execution(self, restart=False):
        pass # pragma: no cover

    # cette operation est appelee par les evenements "Throwing"
    def execute(self):
        pass # pragma: no cover

    def stop(self):
        pass # pragma: no cover

    @property
    def definition(self):
        return self.event.definition.eventKind

class ConditionalEvent(EventKind):

    def validate(self):
        return self.definition.condition(self.event.process)

    def prepare_for_execution(self, restart=False):
        self._push_callback()

    def _callback(self):
        if self.event._p_oid in callbacks:
            del callbacks[self.event._p_oid]

        if self.validate():
            #log.info('%s %s', self.event, "validate ok")
            wi = self.event._get_workitem()
            if wi is not None:
                self.event.start(None)
            else:
                self.stop()
        else:
            #log.info('%s %s', self.event, "validate ko, retry in 1s")
            self._push_callback()

    def _push_callback(self):
        push_event_callback_after_commit(self.event, self._callback, (), 1000)

    def stop(self):
        get_socket().send_pyobj(('stop', self.event._p_oid))

class TerminateEvent(EventKind):
    pass

# https://github.com/zeromq/pyzmq/blob/master/examples/poll/pubsub.py
# http://learning-0mq-with-pyzmq.readthedocs.org/en/latest/pyzmq/multisocket/tornadoeventloop.html

def get_signal_socket_url():
    return 'tcp://127.0.0.1:12346'

class Listener(object):
    def __init__(self, job, identifier):
        self.job = job
        self.identifier = identifier

    def start(self):
        identifier = self.identifier
        job = self.job
        def execute_next(msg):
            # We can't use zodb object from outside here because
            # this code is executed in another thread (eventloop)
            # We don't have site or interaction, so the job must be created
            # before.
            # we can't use push_event_callback_after_commit here because
            # it will never commit in this thread (eventloop)
            if identifier in callbacks:
                callbacks[identifier].close()
                del callbacks[identifier]

            job.args = (msg, )
            # wait 2s that the throw event transaction has committed
            dc = DelayedCallback(job, 2000)
            dc.start()

        ctx = get_zmq_context()
        s = ctx.socket(zmq.SUB)
        s.setsockopt_string(zmq.SUBSCRIBE, u'')
        s.connect(get_signal_socket_url())
        stream = ZMQStream(s)
        callbacks[identifier] = stream
        stream.on_recv(execute_next)

class SignalEvent(EventKind):

    _msg = None

    def validate(self):
        return self.definition.refSignal(self.event.process) == self._msg

    def prepare_for_execution(self, restart=False):
        job = EventJob('system')
        transaction.commit()  # needed to have self.event._p_oid
        job.callable = self._callback
        event_oid = self.event._p_oid
        listener = Listener(job, event_oid)
        get_socket().send_pyobj(('start', listener))

    def stop(self):
        # close ZMQStream
        get_socket().send_pyobj(('close', self.event._p_oid))

    def _callback(self, msg):
        # convert str (py27) / bytes (py34) to unicode (py27) or str (py34)
        self._msg = msg[0].decode('utf-8')
        if self.validate():
            wi = self.event._get_workitem()
            if wi is not None:
                self.event.start(None)
            else:
                self.event.stop()

    def execute(self):
        ref = self.definition.refSignal(self.event.process)
        ctx = get_zmq_context()
        s = ctx.socket(zmq.PUB)
        s.bind(get_signal_socket_url())
        # Sleep to allow sockets to connect.
        time.sleep(0.2)
        s.send_string(ref)
        s.close()

class TimerEvent(EventKind):

    def __init__(self):
        super(TimerEvent, self).__init__()
        self.time_date = None
        self.time_duration = None
        self.time_cycle = None

    def _prepare_time(self, time, restart=False):
        if getattr(self, time, None) is None or not restart:
            setattr(self, time, getattr(self.definition, time)(self.event.process))#TODO

    def _start_time(self, restart=False):
        """Return start time in milliseconds.
        """
        if self.definition.time_date is not None:
            self._prepare_time('time_date', restart)
            if not isinstance(self.time_date, datetime.datetime):
                raise TypeError("Unsupported deadline %r" % self.time_date)
            return (time.mktime(self.time_date.timetuple()) - time.time()) * 1000
        elif self.definition.time_duration is not None:
            self._prepare_time('time_duration', restart)
            if not isinstance(self.time_duration, datetime.timedelta):
                raise TypeError("Unsupported deadline %r" % self.time_duration)
            return self.time_duration.total_seconds() * 1000
        elif self.definition.time_cycle is not None:
            raise NotImplementedError
        else:
            raise TypeError("Unsupported deadline")

    def validate(self):
        return True

    def prepare_for_execution(self, restart=False):
        self._push_callback(restart)

    def _callback(self):
        if self.event._p_oid in callbacks:
            del callbacks[self.event._p_oid]

        if self.validate():
            wi = self.event._get_workitem()
            if wi is not None:
                self.event.start(None)
            else:
                self.event.stop()
        else:
            self._push_callback()

    def _push_callback(self, restart=False):
        deadline = self._start_time(restart)
        push_event_callback_after_commit(self.event, self._callback, (), deadline)

    def stop(self):
        # stop DelayedCallback
        get_socket().send_pyobj(('stop', self.event._p_oid))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# Dactylo -- A datasets activity streams logger
# By: Emmanuel Raviart <emmanuel@raviart.com>
#
# Copyright (C) 2013 Etalab
# http://github.com/etalab/dactylo
#
# This file is part of Dactylo.
#
# Dactylo is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# Dactylo is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""Controllers for states"""

import collections
import json
import logging

from .. import contexts, conv, model, urls, wsgihelpers

json_to_metrics = conv.pipe(
    conv.test_isinstance(dict),
#    conv.struct(
#        dict(
#            datasets = conv.pipe(
#                conv.test_isinstance(dict),
#                conv.struct(
#                    dict(
#                        count = conv.pipe(
#                            conv.test_isinstance(int),
#                            conv.test_greater_or_equal(0),
#                            conv.not_none,
#                            ),
#                        weights = conv.pipe(
#                            conv.test_isinstance(float),
#                            conv.test_greater_or_equal(0.0),
#                            conv.not_none,
#                            ),
#                        ),
#                    ),
#                conv.not_none,
#                ),
#            ),
#        ),
    )
log = logging.getLogger(__name__)

@wsgihelpers.wsgify
def api1_get(req):
    ctx = contexts.Ctx(req)
    headers = wsgihelpers.handle_cross_origin_resource_sharing(ctx)

    assert req.method == 'GET', req.method
    params = req.GET
    inputs = dict(
        callback = params.get('callback'),
        context = params.get('context'),
        )
    data, errors = conv.pipe(
        conv.struct(
            dict(
                callback = conv.pipe(
                    conv.test_isinstance(basestring),
                    conv.cleanup_line,
                    ),
                context = conv.test_isinstance(basestring),
                ),
            ),
        )(inputs, state = ctx)
    if errors is not None:
        return wsgihelpers.respond_json(ctx,
            dict(
                apiVersion = '1.0',
                context = inputs['context'],
                error = dict(
                    code = 400,  # Bad Request
                    errors = [
                        dict(
                            location = key,
                            message = error,
                            )
                        for key, error in sorted(errors.iteritems())
                        ],
                    # message will be automatically defined.
                    ),
                method = req.script_name,
                params = inputs,
                url = req.url.decode('utf-8'),
                ),
            headers = headers,
            jsonp = inputs['callback'],
            )

    return wsgihelpers.respond_json(ctx,
        collections.OrderedDict(sorted(dict(
            apiVersion = '1.0',
            context = data['context'],
            method = req.script_name,
            params = inputs,
            url = req.url.decode('utf-8'),
            value = model.metrics,
            ).iteritems())),
        headers = headers,
        jsonp = data['callback'],
        )

@wsgihelpers.wsgify
def api1_set(req):
    ctx = contexts.Ctx(req)
    headers = wsgihelpers.handle_cross_origin_resource_sharing(ctx)

    assert req.method == 'POST', req.method

    inputs_converters = dict(
        # Shared secret between client and server
        api_key = conv.pipe(
            conv.test_isinstance(basestring),
            conv.input_to_token,
            conv.not_none,
            ),
        # For asynchronous calls
        context = conv.test_isinstance(basestring),
        # "value" is handled below.
        )

    content_type = req.content_type
    if content_type is not None:
        content_type = content_type.split(';', 1)[0].strip()
    if content_type == 'application/json':
        inputs, error = conv.pipe(
            conv.make_input_to_json(),
            conv.test_isinstance(dict),
            )(req.body, state = ctx)
        if error is not None:
            return wsgihelpers.respond_json(ctx,
                collections.OrderedDict(sorted(dict(
                    apiVersion = '1.0',
                    error = collections.OrderedDict(sorted(dict(
                        code = 400,  # Bad Request
                        errors = [error],
                        message = ctx._(u'Invalid JSON in request POST body'),
                        ).iteritems())),
                    method = req.script_name,
                    params = req.body,
                    url = req.url.decode('utf-8'),
                    ).iteritems())),
                headers = headers,
                )
        inputs_converters.update(dict(
            value = conv.pipe(
                json_to_metrics,
                conv.not_none,
                ),
            ))
    else:
        # URL-encoded POST.
        inputs = dict(req.POST)
        inputs_converters.update(dict(
            value = conv.pipe(
                conv.make_input_to_json(),
                json_to_metrics,
                conv.not_none,
                ),
            ))

    data, errors = conv.struct(inputs_converters)(inputs, state = ctx)
    if errors is not None:
        return wsgihelpers.respond_json(ctx,
            collections.OrderedDict(sorted(dict(
                apiVersion = '1.0',
                context = inputs.get('context'),
                error = collections.OrderedDict(sorted(dict(
                    code = 400,  # Bad Request
                    errors = [errors],
                    message = ctx._(u'Bad parameters in request'),
                    ).iteritems())),
                method = req.script_name,
                params = inputs,
                url = req.url.decode('utf-8'),
                ).iteritems())),
            headers = headers,
            )

    api_key = data['api_key']
    account = model.Account.find_one(
        dict(
            api_key = api_key,
            ),
        as_class = collections.OrderedDict,
        )
    if account is None:
        return wsgihelpers.respond_json(ctx,
            collections.OrderedDict(sorted(dict(
                apiVersion = '1.0',
                context = data['context'],
                error = collections.OrderedDict(sorted(dict(
                    code = 401,  # Unauthorized
                    message = ctx._('Unknown API Key: {}').format(api_key),
                    ).iteritems())),
                method = req.script_name,
                params = inputs,
                url = req.url.decode('utf-8'),
                ).iteritems())),
            headers = headers,
            )

    model.metrics = data['value']
    message = unicode(json.dumps(model.metrics, encoding = 'utf-8', ensure_ascii = False, indent = 2))
    for client in model.websocket_metrics_clients:
        client.send(message)

    return wsgihelpers.respond_json(ctx,
        collections.OrderedDict(sorted(dict(
            apiVersion = '1.0',
            context = data['context'],
            method = req.script_name,
            params = inputs,
            url = req.url.decode('utf-8'),
            value = model.metrics,
            ).iteritems())),
        headers = headers,
        )

def route_api1_class(environ, start_response):
    router = urls.make_router(
        ('GET', '^/?$', api1_get),
        ('POST', '^/?$', api1_set),
        )
    return router(environ, start_response)
<commit_msg>Add test controller.<commit_after># -*- coding: utf-8 -*-

# Dactylo -- A datasets activity streams logger
# By: Emmanuel Raviart <emmanuel@raviart.com>
#
# Copyright (C) 2013 Etalab
# http://github.com/etalab/dactylo
#
# This file is part of Dactylo.
#
# Dactylo is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# Dactylo is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""Controllers for states"""

import collections
import datetime
import json
import logging

from .. import contexts, conv, model, urls, wsgihelpers

json_to_metrics = conv.pipe(
    conv.test_isinstance(dict),
#    conv.struct(
#        dict(
#            datasets = conv.pipe(
#                conv.test_isinstance(dict),
#                conv.struct(
#                    dict(
#                        count = conv.pipe(
#                            conv.test_isinstance(int),
#                            conv.test_greater_or_equal(0),
#                            conv.not_none,
#                            ),
#                        weights = conv.pipe(
#                            conv.test_isinstance(float),
#                            conv.test_greater_or_equal(0.0),
#                            conv.not_none,
#                            ),
#                        ),
#                    ),
#                conv.not_none,
#                ),
#            ),
#        ),
    )
log = logging.getLogger(__name__)

@wsgihelpers.wsgify
def api1_get(req):
    ctx = contexts.Ctx(req)
    headers = wsgihelpers.handle_cross_origin_resource_sharing(ctx)

    assert req.method == 'GET', req.method
    params = req.GET
    inputs = dict(
        callback = params.get('callback'),
        context = params.get('context'),
        )
    data, errors = conv.pipe(
        conv.struct(
            dict(
                callback = conv.pipe(
                    conv.test_isinstance(basestring),
                    conv.cleanup_line,
                    ),
                context = conv.test_isinstance(basestring),
                ),
            ),
        )(inputs, state = ctx)
    if errors is not None:
        return wsgihelpers.respond_json(ctx,
            dict(
                apiVersion = '1.0',
                context = inputs['context'],
                error = dict(
                    code = 400,  # Bad Request
                    errors = [
                        dict(
                            location = key,
                            message = error,
                            )
                        for key, error in sorted(errors.iteritems())
                        ],
                    # message will be automatically defined.
                    ),
                method = req.script_name,
                params = inputs,
                url = req.url.decode('utf-8'),
                ),
            headers = headers,
            jsonp = inputs['callback'],
            )

    return wsgihelpers.respond_json(ctx,
        collections.OrderedDict(sorted(dict(
            apiVersion = '1.0',
            context = data['context'],
            method = req.script_name,
            params = inputs,
            url = req.url.decode('utf-8'),
            value = model.metrics,
            ).iteritems())),
        headers = headers,
        jsonp = data['callback'],
        )

@wsgihelpers.wsgify
def api1_set(req):
    ctx = contexts.Ctx(req)
    headers = wsgihelpers.handle_cross_origin_resource_sharing(ctx)

    assert req.method == 'POST', req.method

    inputs_converters = dict(
        # Shared secret between client and server
        api_key = conv.pipe(
            conv.test_isinstance(basestring),
            conv.input_to_token,
            conv.not_none,
            ),
        # For asynchronous calls
        context = conv.test_isinstance(basestring),
        # "value" is handled below.
        )

    content_type = req.content_type
    if content_type is not None:
        content_type = content_type.split(';', 1)[0].strip()
    if content_type == 'application/json':
        inputs, error = conv.pipe(
            conv.make_input_to_json(),
            conv.test_isinstance(dict),
            )(req.body, state = ctx)
        if error is not None:
            return wsgihelpers.respond_json(ctx,
                collections.OrderedDict(sorted(dict(
                    apiVersion = '1.0',
                    error = collections.OrderedDict(sorted(dict(
                        code = 400,  # Bad Request
                        errors = [error],
                        message = ctx._(u'Invalid JSON in request POST body'),
                        ).iteritems())),
                    method = req.script_name,
                    params = req.body,
                    url = req.url.decode('utf-8'),
                    ).iteritems())),
                headers = headers,
                )
        inputs_converters.update(dict(
            value = conv.pipe(
                json_to_metrics,
                conv.not_none,
                ),
            ))
    else:
        # URL-encoded POST.
        inputs = dict(req.POST)
        inputs_converters.update(dict(
            value = conv.pipe(
                conv.make_input_to_json(),
                json_to_metrics,
                conv.not_none,
                ),
            ))

    data, errors = conv.struct(inputs_converters)(inputs, state = ctx)
    if errors is not None:
        return wsgihelpers.respond_json(ctx,
            collections.OrderedDict(sorted(dict(
                apiVersion = '1.0',
                context = inputs.get('context'),
                error = collections.OrderedDict(sorted(dict(
                    code = 400,  # Bad Request
                    errors = [errors],
                    message = ctx._(u'Bad parameters in request'),
                    ).iteritems())),
                method = req.script_name,
                params = inputs,
                url = req.url.decode('utf-8'),
                ).iteritems())),
            headers = headers,
            )

    api_key = data['api_key']
    account = model.Account.find_one(
        dict(
            api_key = api_key,
            ),
        as_class = collections.OrderedDict,
        )
    if account is None:
        return wsgihelpers.respond_json(ctx,
            collections.OrderedDict(sorted(dict(
                apiVersion = '1.0',
                context = data['context'],
                error = collections.OrderedDict(sorted(dict(
                    code = 401,  # Unauthorized
                    message = ctx._('Unknown API Key: {}').format(api_key),
                    ).iteritems())),
                method = req.script_name,
                params = inputs,
                url = req.url.decode('utf-8'),
                ).iteritems())),
            headers = headers,
            )

    model.metrics = data['value']
    message = unicode(json.dumps(model.metrics, encoding = 'utf-8', ensure_ascii = False, indent = 2))
    for client in model.websocket_metrics_clients:
        client.send(message)

    return wsgihelpers.respond_json(ctx,
        collections.OrderedDict(sorted(dict(
            apiVersion = '1.0',
            context = data['context'],
            method = req.script_name,
            params = inputs,
            url = req.url.decode('utf-8'),
            value = model.metrics,
            ).iteritems())),
        headers = headers,
        )

@wsgihelpers.wsgify
def api1_test(req):
    ctx = contexts.Ctx(req)
    headers = wsgihelpers.handle_cross_origin_resource_sharing(ctx)

    assert req.method == 'GET', req.method
    params = req.GET
    inputs = dict(
        callback = params.get('callback'),
        context = params.get('context'),
        )
    data, errors = conv.pipe(
        conv.struct(
            dict(
                callback = conv.pipe(
                    conv.test_isinstance(basestring),
                    conv.cleanup_line,
                    ),
                context = conv.test_isinstance(basestring),
                ),
            ),
        )(inputs, state = ctx)
    if errors is not None:
        return wsgihelpers.respond_json(ctx,
            dict(
                apiVersion = '1.0',
                context = inputs['context'],
                error = dict(
                    code = 400,  # Bad Request
                    errors = [
                        dict(
                            location = key,
                            message = error,
                            )
                        for key, error in sorted(errors.iteritems())
                        ],
                    # message will be automatically defined.
                    ),
                method = req.script_name,
                params = inputs,
                url = req.url.decode('utf-8'),
                ),
            headers = headers,
            jsonp = inputs['callback'],
            )

    test_metrics = model.metrics.copy()
    test_metrics['test'] = datetime.datetime.utcnow().isoformat(),
    message = unicode(json.dumps(test_metrics, encoding = 'utf-8', ensure_ascii = False, indent = 2))
    for client in model.websocket_metrics_clients:
        client.send(message)

    return wsgihelpers.respond_json(ctx,
        collections.OrderedDict(sorted(dict(
            apiVersion = '1.0',
            context = data['context'],
            method = req.script_name,
            params = inputs,
            url = req.url.decode('utf-8'),
            value = model.metrics,
            ).iteritems())),
        headers = headers,
        jsonp = data['callback'],
        )

def route_api1_class(environ, start_response):
    router = urls.make_router(
        ('GET', '^/?$', api1_get),
        ('GET', '^/test/?$', api1_test),
        ('POST', '^/?$', api1_set),
        )
    return router(environ, start_response)
<|endoftext|>


<commit_before>import os
import posixpath

import pytest
from tlz import concat
from fsspec.core import open_files, get_fs_token_paths

import dask
import dask.bag as db
from dask.bytes.core import read_bytes

try:
    import distributed
    from distributed import Client
    from distributed.utils_test import cluster, loop  # noqa: F401
except (ImportError, SyntaxError):
    distributed = None

try:
    import pyarrow
except ImportError:
    pyarrow = None

if not os.environ.get("DASK_RUN_HDFS_TESTS", ""):
    pytestmark = pytest.mark.skip(reason="HDFS tests not configured to run")

basedir = "/tmp/test-dask"

@pytest.fixture
def hdfs(request):
    hdfs = pyarrow.hdfs.connect(host="localhost", port=8020)

    if hdfs.exists(basedir):
        hdfs.rm(basedir, recursive=True)
    hdfs.mkdir(basedir)

    yield hdfs

    if hdfs.exists(basedir):
        hdfs.rm(basedir, recursive=True)

# This mark doesn't check the minimum pyarrow version.
require_pyarrow = pytest.mark.skipif(not pyarrow, reason="pyarrow not installed")

def test_read_bytes(hdfs):
    nfiles = 10

    data = b"a" * int(1e3)

    for fn in ["%s/file.%d" % (basedir, i) for i in range(nfiles)]:
        with hdfs.open(fn, "wb", replication=1) as f:
            f.write(data)

    sample, values = read_bytes("hdfs://%s/file.*" % basedir)

    (results,) = dask.compute(values)
    assert [b"".join(r) for r in results] == nfiles * [data]

def test_read_bytes_URL(hdfs):
    nfiles = 10
    data = b"a" * int(1e3)

    for fn in ["%s/file.%d" % (basedir, i) for i in range(nfiles)]:
        with hdfs.open(fn, "wb", replication=1) as f:
            f.write(data)

    path = "hdfs://localhost:8020%s/file.*" % basedir
    sample, values = read_bytes(path)

    (results,) = dask.compute(values)
    assert [b"".join(r) for r in results] == nfiles * [data]

def test_read_bytes_big_file(hdfs):
    fn = "%s/file" % basedir

    # Write 100 MB file
    nblocks = int(1e3)
    blocksize = int(1e5)
    data = b"a" * blocksize
    with hdfs.open(fn, "wb", replication=1) as f:
        for i in range(nblocks):
            f.write(data)

    sample, values = read_bytes("hdfs://" + fn, blocksize=blocksize)

    assert sample[:5] == b"aaaaa"
    assert len(values[0]) == nblocks

    (results,) = dask.compute(values[0])
    assert sum(map(len, results)) == nblocks * blocksize
    for r in results:
        assert set(r.decode("utf-8")) == {"a"}

def test_deterministic_key_names(hdfs):
    data = b"abc\n" * int(1e3)
    fn = "%s/file" % basedir

    with hdfs.open(fn, "wb", replication=1) as fil:
        fil.write(data)

    _, x = read_bytes("hdfs://%s/*" % basedir, delimiter=b"\n", sample=False)
    _, y = read_bytes("hdfs://%s/*" % basedir, delimiter=b"\n", sample=False)
    _, z = read_bytes("hdfs://%s/*" % basedir, delimiter=b"c", sample=False)

    assert [f.key for f in concat(x)] == [f.key for f in concat(y)]
    assert [f.key for f in concat(x)] != [f.key for f in concat(z)]

def test_open_files_write(hdfs):
    path = "hdfs://%s/" % basedir
    data = [b"test data %i" % i for i in range(5)]

    files = open_files(path, num=len(data), mode="wb")
    for fil, b in zip(files, data):
        with fil as f:
            f.write(b)

    sample, vals = read_bytes("hdfs://%s/*.part" % basedir)

    (results,) = dask.compute(list(concat(vals)))
    assert data == results

def test_read_csv(hdfs):
    dd = pytest.importorskip("dask.dataframe")

    with hdfs.open("%s/1.csv" % basedir, "wb") as f:
        f.write(b"name,amount,id\nAlice,100,1\nBob,200,2")

    with hdfs.open("%s/2.csv" % basedir, "wb") as f:
        f.write(b"name,amount,id\nCharlie,300,3\nDennis,400,4")

    df = dd.read_csv("hdfs://%s/*.csv" % basedir)

    assert isinstance(df, dd.DataFrame)
    assert df.id.sum().compute() == 1 + 2 + 3 + 4

def test_read_text(hdfs):
    import multiprocessing as mp
    from concurrent.futures import ProcessPoolExecutor

    ctx = mp.get_context("spawn")

    with ProcessPoolExecutor(2, ctx) as pool:
        with hdfs.open("%s/text.1.txt" % basedir, "wb") as f:
            f.write("Alice 100\nBob 200\nCharlie 300".encode())

        with hdfs.open("%s/text.2.txt" % basedir, "wb") as f:
            f.write("Dan 400\nEdith 500\nFrank 600".encode())

        with hdfs.open("%s/other.txt" % basedir, "wb") as f:
            f.write("a b\nc d".encode())

        b = db.read_text("hdfs://%s/text.*.txt" % basedir)
        with dask.config.set(pool=pool):
            result = b.str.strip().str.split().map(len).compute()

        assert result == [2, 2, 2, 2, 2, 2]

        b = db.read_text("hdfs://%s/other.txt" % basedir)
        with dask.config.set(pool=pool):
            result = b.str.split().flatten().compute()

        assert result == ["a", "b", "c", "d"]

def test_read_text_unicode(hdfs):
    data = b"abcd\xc3\xa9"
    fn = "%s/data.txt" % basedir
    with hdfs.open(fn, "wb") as f:
        f.write(b"\n".join([data, data]))

    f = db.read_text("hdfs://" + fn, collection=False)

    result = f[0].compute()
    assert len(result) == 2
    assert list(map(str.strip, result)) == [data.decode("utf-8")] * 2
    assert len(result[0].strip()) == 5

@require_pyarrow
def test_parquet_pyarrow(hdfs):
    dd = pytest.importorskip("dask.dataframe")
    import pandas as pd
    import numpy as np

    fn = "%s/test.parquet" % basedir
    hdfs_fn = "hdfs://%s" % fn
    df = pd.DataFrame(np.random.normal(size=(1000, 4)), columns=list("abcd"))
    ddf = dd.from_pandas(df, npartitions=4)

    ddf.to_parquet(hdfs_fn, engine="pyarrow")

    assert len(hdfs.ls(fn))  # Files are written

    ddf2 = dd.read_parquet(hdfs_fn, engine="pyarrow")

    assert len(ddf2) == 1000  # smoke test on read

def test_glob(hdfs):

    tree = {
        basedir: (["c", "c2"], ["a", "a1", "a2", "a3", "b1"]),
        basedir + "/c": (["d"], ["x1", "x2"]),
        basedir + "/c2": (["d"], ["x1", "x2"]),
        basedir + "/c/d": ([], ["x3"]),
    }

    hdfs, _, _ = get_fs_token_paths("hdfs:///")
    hdfs.makedirs(basedir + "/c/d")
    hdfs.makedirs(basedir + "/c2/d/")
    for fn in (
        posixpath.join(dirname, f)
        for (dirname, (_, fils)) in tree.items()
        for f in fils
    ):
        with hdfs.open(fn, mode="wb") as f2:
            f2.write(b"000")

    assert set(hdfs.glob(basedir + "/a*")) == {
        basedir + p for p in ["/a", "/a1", "/a2", "/a3"]
    }

    assert set(hdfs.glob(basedir + "/c/*")) == {
        basedir + p for p in ["/c/x1", "/c/x2", "/c/d"]
    }

    assert set(hdfs.glob(basedir + "/*/x*")) == {
        basedir + p for p in ["/c/x1", "/c/x2", "/c2/x1", "/c2/x2"]
    }
    assert set(hdfs.glob(basedir + "/*/x1")) == {
        basedir + p for p in ["/c/x1", "/c2/x1"]
    }

    assert hdfs.find("/this-path-doesnt-exist") == []
    assert hdfs.find(basedir + "/missing/") == []
    assert hdfs.find(basedir + "/missing/x1") == []
    assert hdfs.glob(basedir + "/missing/*") == []
    assert hdfs.glob(basedir + "/*/missing") == []

    assert set(hdfs.glob(basedir + "/*")) == {
        basedir + p for p in ["/a", "/a1", "/a2", "/a3", "/b1", "/c", "/c2"]
    }

@pytest.mark.skipif(
    not distributed, reason="Skipped as distributed is not installed."  # noqa: F811
)  # noqa: F811
def test_distributed(hdfs, loop):  # noqa: F811
    dd = pytest.importorskip("dask.dataframe")

    with hdfs.open("%s/1.csv" % basedir, "wb") as f:
        f.write(b"name,amount,id\nAlice,100,1\nBob,200,2")

    with hdfs.open("%s/2.csv" % basedir, "wb") as f:
        f.write(b"name,amount,id\nCharlie,300,3\nDennis,400,4")

    with cluster() as (s, [a, b]):
        with Client(s["address"], loop=loop):  # noqa: F811
            df = dd.read_csv("hdfs://%s/*.csv" % basedir)
            assert df.id.sum().compute() == 1 + 2 + 3 + 4
<commit_msg>Use `multiprocessing.Pool` in `test_read_text` (#7472)<commit_after>import os
import posixpath

import pytest
from tlz import concat
from fsspec.core import open_files, get_fs_token_paths

import dask
import dask.bag as db
from dask.bytes.core import read_bytes

try:
    import distributed
    from distributed import Client
    from distributed.utils_test import cluster, loop  # noqa: F401
except (ImportError, SyntaxError):
    distributed = None

try:
    import pyarrow
except ImportError:
    pyarrow = None

if not os.environ.get("DASK_RUN_HDFS_TESTS", ""):
    pytestmark = pytest.mark.skip(reason="HDFS tests not configured to run")

basedir = "/tmp/test-dask"

@pytest.fixture
def hdfs(request):
    hdfs = pyarrow.hdfs.connect(host="localhost", port=8020)

    if hdfs.exists(basedir):
        hdfs.rm(basedir, recursive=True)
    hdfs.mkdir(basedir)

    yield hdfs

    if hdfs.exists(basedir):
        hdfs.rm(basedir, recursive=True)

# This mark doesn't check the minimum pyarrow version.
require_pyarrow = pytest.mark.skipif(not pyarrow, reason="pyarrow not installed")

def test_read_bytes(hdfs):
    nfiles = 10

    data = b"a" * int(1e3)

    for fn in ["%s/file.%d" % (basedir, i) for i in range(nfiles)]:
        with hdfs.open(fn, "wb", replication=1) as f:
            f.write(data)

    sample, values = read_bytes("hdfs://%s/file.*" % basedir)

    (results,) = dask.compute(values)
    assert [b"".join(r) for r in results] == nfiles * [data]

def test_read_bytes_URL(hdfs):
    nfiles = 10
    data = b"a" * int(1e3)

    for fn in ["%s/file.%d" % (basedir, i) for i in range(nfiles)]:
        with hdfs.open(fn, "wb", replication=1) as f:
            f.write(data)

    path = "hdfs://localhost:8020%s/file.*" % basedir
    sample, values = read_bytes(path)

    (results,) = dask.compute(values)
    assert [b"".join(r) for r in results] == nfiles * [data]

def test_read_bytes_big_file(hdfs):
    fn = "%s/file" % basedir

    # Write 100 MB file
    nblocks = int(1e3)
    blocksize = int(1e5)
    data = b"a" * blocksize
    with hdfs.open(fn, "wb", replication=1) as f:
        for i in range(nblocks):
            f.write(data)

    sample, values = read_bytes("hdfs://" + fn, blocksize=blocksize)

    assert sample[:5] == b"aaaaa"
    assert len(values[0]) == nblocks

    (results,) = dask.compute(values[0])
    assert sum(map(len, results)) == nblocks * blocksize
    for r in results:
        assert set(r.decode("utf-8")) == {"a"}

def test_deterministic_key_names(hdfs):
    data = b"abc\n" * int(1e3)
    fn = "%s/file" % basedir

    with hdfs.open(fn, "wb", replication=1) as fil:
        fil.write(data)

    _, x = read_bytes("hdfs://%s/*" % basedir, delimiter=b"\n", sample=False)
    _, y = read_bytes("hdfs://%s/*" % basedir, delimiter=b"\n", sample=False)
    _, z = read_bytes("hdfs://%s/*" % basedir, delimiter=b"c", sample=False)

    assert [f.key for f in concat(x)] == [f.key for f in concat(y)]
    assert [f.key for f in concat(x)] != [f.key for f in concat(z)]

def test_open_files_write(hdfs):
    path = "hdfs://%s/" % basedir
    data = [b"test data %i" % i for i in range(5)]

    files = open_files(path, num=len(data), mode="wb")
    for fil, b in zip(files, data):
        with fil as f:
            f.write(b)

    sample, vals = read_bytes("hdfs://%s/*.part" % basedir)

    (results,) = dask.compute(list(concat(vals)))
    assert data == results

def test_read_csv(hdfs):
    dd = pytest.importorskip("dask.dataframe")

    with hdfs.open("%s/1.csv" % basedir, "wb") as f:
        f.write(b"name,amount,id\nAlice,100,1\nBob,200,2")

    with hdfs.open("%s/2.csv" % basedir, "wb") as f:
        f.write(b"name,amount,id\nCharlie,300,3\nDennis,400,4")

    df = dd.read_csv("hdfs://%s/*.csv" % basedir)

    assert isinstance(df, dd.DataFrame)
    assert df.id.sum().compute() == 1 + 2 + 3 + 4

def test_read_text(hdfs):
    import multiprocessing as mp

    pool = mp.get_context("spawn").Pool(2)

    with pool:
        with hdfs.open("%s/text.1.txt" % basedir, "wb") as f:
            f.write("Alice 100\nBob 200\nCharlie 300".encode())

        with hdfs.open("%s/text.2.txt" % basedir, "wb") as f:
            f.write("Dan 400\nEdith 500\nFrank 600".encode())

        with hdfs.open("%s/other.txt" % basedir, "wb") as f:
            f.write("a b\nc d".encode())

        b = db.read_text("hdfs://%s/text.*.txt" % basedir)
        with dask.config.set(pool=pool):
            result = b.str.strip().str.split().map(len).compute()

        assert result == [2, 2, 2, 2, 2, 2]

        b = db.read_text("hdfs://%s/other.txt" % basedir)
        with dask.config.set(pool=pool):
            result = b.str.split().flatten().compute()

        assert result == ["a", "b", "c", "d"]

def test_read_text_unicode(hdfs):
    data = b"abcd\xc3\xa9"
    fn = "%s/data.txt" % basedir
    with hdfs.open(fn, "wb") as f:
        f.write(b"\n".join([data, data]))

    f = db.read_text("hdfs://" + fn, collection=False)

    result = f[0].compute()
    assert len(result) == 2
    assert list(map(str.strip, result)) == [data.decode("utf-8")] * 2
    assert len(result[0].strip()) == 5

@require_pyarrow
def test_parquet_pyarrow(hdfs):
    dd = pytest.importorskip("dask.dataframe")
    import pandas as pd
    import numpy as np

    fn = "%s/test.parquet" % basedir
    hdfs_fn = "hdfs://%s" % fn
    df = pd.DataFrame(np.random.normal(size=(1000, 4)), columns=list("abcd"))
    ddf = dd.from_pandas(df, npartitions=4)

    ddf.to_parquet(hdfs_fn, engine="pyarrow")

    assert len(hdfs.ls(fn))  # Files are written

    ddf2 = dd.read_parquet(hdfs_fn, engine="pyarrow")

    assert len(ddf2) == 1000  # smoke test on read

def test_glob(hdfs):

    tree = {
        basedir: (["c", "c2"], ["a", "a1", "a2", "a3", "b1"]),
        basedir + "/c": (["d"], ["x1", "x2"]),
        basedir + "/c2": (["d"], ["x1", "x2"]),
        basedir + "/c/d": ([], ["x3"]),
    }

    hdfs, _, _ = get_fs_token_paths("hdfs:///")
    hdfs.makedirs(basedir + "/c/d")
    hdfs.makedirs(basedir + "/c2/d/")
    for fn in (
        posixpath.join(dirname, f)
        for (dirname, (_, fils)) in tree.items()
        for f in fils
    ):
        with hdfs.open(fn, mode="wb") as f2:
            f2.write(b"000")

    assert set(hdfs.glob(basedir + "/a*")) == {
        basedir + p for p in ["/a", "/a1", "/a2", "/a3"]
    }

    assert set(hdfs.glob(basedir + "/c/*")) == {
        basedir + p for p in ["/c/x1", "/c/x2", "/c/d"]
    }

    assert set(hdfs.glob(basedir + "/*/x*")) == {
        basedir + p for p in ["/c/x1", "/c/x2", "/c2/x1", "/c2/x2"]
    }
    assert set(hdfs.glob(basedir + "/*/x1")) == {
        basedir + p for p in ["/c/x1", "/c2/x1"]
    }

    assert hdfs.find("/this-path-doesnt-exist") == []
    assert hdfs.find(basedir + "/missing/") == []
    assert hdfs.find(basedir + "/missing/x1") == []
    assert hdfs.glob(basedir + "/missing/*") == []
    assert hdfs.glob(basedir + "/*/missing") == []

    assert set(hdfs.glob(basedir + "/*")) == {
        basedir + p for p in ["/a", "/a1", "/a2", "/a3", "/b1", "/c", "/c2"]
    }

@pytest.mark.skipif(
    not distributed, reason="Skipped as distributed is not installed."  # noqa: F811
)  # noqa: F811
def test_distributed(hdfs, loop):  # noqa: F811
    dd = pytest.importorskip("dask.dataframe")

    with hdfs.open("%s/1.csv" % basedir, "wb") as f:
        f.write(b"name,amount,id\nAlice,100,1\nBob,200,2")

    with hdfs.open("%s/2.csv" % basedir, "wb") as f:
        f.write(b"name,amount,id\nCharlie,300,3\nDennis,400,4")

    with cluster() as (s, [a, b]):
        with Client(s["address"], loop=loop):  # noqa: F811
            df = dd.read_csv("hdfs://%s/*.csv" % basedir)
            assert df.id.sum().compute() == 1 + 2 + 3 + 4
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os, sys
from setuptools import setup

if sys.argv[-1] == 'publish':
    os.system('python setup.py sdist upload')
    sys.exit()

readme = open('README.rst').read()
long_description = readme
doclink = '''
Documentation
-------------

The full documentation is at http://workbench.rtfd.org. '''

exec(open('workbench_cli/version.py').read())
setup(
    name='workbench_cli',
    version=__version__,
    description='Command Line Interface for Workbench',
    long_description=readme + '\n\n' + doclink + '\n\n',
    author='The Workbench Team',
    author_email='support@supercowpowers.com',
    url='http://github.com/SuperCowPowers/workbench',
    packages=['workbench_cli'],
    package_dir={'workbench_cli': 'workbench_cli'},
    include_package_data=True,
    scripts=['workbench_cli/workbench'],
    install_requires=['colorama', 'ipython', 'lz4', 'pandas', 'pytest', 'zerorpc'],
    license='MIT',
    zip_safe=False,
    keywords='workbench security python',
    classifiers=[
        'Development Status :: 2 - Pre-Alpha',
        'Intended Audience :: Developers',
        'License :: OSI Approved :: MIT License',
        'Natural Language :: English',
        'Programming Language :: Python :: 2.7',
    ]
)
<commit_msg>adding matplotlib<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os, sys
from setuptools import setup

if sys.argv[-1] == 'publish':
    os.system('python setup.py sdist upload')
    sys.exit()

readme = open('README.rst').read()
long_description = readme
doclink = '''
Documentation
-------------

The full documentation is at http://workbench.rtfd.org. '''

exec(open('workbench_cli/version.py').read())
setup(
    name='workbench_cli',
    version=__version__,
    description='Command Line Interface for Workbench',
    long_description=readme + '\n\n' + doclink + '\n\n',
    author='The Workbench Team',
    author_email='support@supercowpowers.com',
    url='http://github.com/SuperCowPowers/workbench',
    packages=['workbench_cli'],
    package_dir={'workbench_cli': 'workbench_cli'},
    include_package_data=True,
    scripts=['workbench_cli/workbench'],
    install_requires=['colorama', 'ipython', 'lz4', 'matplotlib',
                      'pandas', 'pytest', 'zerorpc'],
    license='MIT',
    zip_safe=False,
    keywords='workbench security python',
    classifiers=[
        'Development Status :: 2 - Pre-Alpha',
        'Intended Audience :: Developers',
        'License :: OSI Approved :: MIT License',
        'Natural Language :: English',
        'Programming Language :: Python :: 2.7',
    ]
)
<|endoftext|>