<commit_before># Copyright 2018 Google Inc.
#
# 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.
"""Colab-specific Google Drive integration."""

from __future__ import absolute_import as _
from __future__ import division as _
from __future__ import print_function as _

import collections as _collections
import os as _os
import signal as _signal
import socket as _socket
import subprocess as _subprocess
import sys as _sys
import tempfile as _tempfile
import uuid as _uuid

import pexpect.popen_spawn as _popen_spawn

from google.colab import output as _output

__all__ = ['flush_and_unmount', 'mount']

_Environment = _collections.namedtuple(
    '_Environment',
    ('home', 'root_dir', 'inet_family', 'dev', 'path', 'config_dir'))

def _env():
  """Create and return an _Environment to use."""
  home = _os.environ['HOME']
  root_dir = _os.path.realpath(
      _os.path.join(_os.environ['CLOUDSDK_CONFIG'], '../..'))
  inet_family = 'IPV4_ONLY'
  dev = '/dev/fuse'
  path = '/usr/local/bin:/usr/local/sbin:/usr/bin:/usr/sbin:/bin:/sbin:.'
  if len(root_dir) > 1 and not root_dir.startswith('/usr/local/google/'):
    home = _os.path.join(root_dir, home)
    inet_family = 'IPV6_ONLY'
    fum = _os.environ['HOME'].split('mount')[0] + '/mount/alloc/fusermount'
    if 'BORG_ALLOC_DIR' in _os.environ:
      # For support in guitar cluster invocations.
      fum = _os.environ['BORG_ALLOC_DIR'] + '/alloc/fusermount'
    dev = fum + '/dev/fuse'
    path = path + ':' + fum + '/bin'
  config_dir = _os.path.join(home, '.config', 'Google')
  return _Environment(
      home=home,
      root_dir=root_dir,
      inet_family=inet_family,
      dev=dev,
      path=path,
      config_dir=config_dir)

def _logs_dir():
  return _os.path.join(_env().config_dir, 'DriveFS/Logs/')

def _timeouts_path():
  return _os.path.join(_logs_dir(), 'timeouts.txt')

def flush_and_unmount(timeout_ms=24 * 60 * 60 * 1000):
  """Unmount Google Drive and flush any outstanding writes to it."""
  env = _env()
  if b'type fuse.drive' not in _subprocess.check_output(['/bin/mount']):
    print('Drive not mounted, so nothing to flush and unmount.')
    return
  drive_bin = _os.path.join(env.root_dir, 'opt/google/drive/drive')
  p = _subprocess.Popen([
      drive_bin, '--push_changes_and_quit', '--single_process',
      '--timeout_sec={}'.format(int(timeout_ms / 1000))
  ],
                        stdout=_subprocess.PIPE,
                        stderr=_subprocess.PIPE)
  out, err = p.communicate()
  if mount._DEBUG:  # pylint:disable=protected-access
    print('flush_and_unmount: out: {}\nerr: {}'.format(out, err))
  if p.returncode:
    raise ValueError('flush_and_unmount failed')

def mount(mountpoint,
          force_remount=False,
          timeout_ms=60000,
          use_metadata_server=False):
  """Mount your Google Drive at the specified mountpoint path."""

  if ' ' in mountpoint:
    raise ValueError('Mountpoint must not contain a space.')

  mountpoint = _os.path.expanduser(mountpoint)
  # If we've already mounted drive at the specified mountpoint, exit now.
  already_mounted = _os.path.isdir(_os.path.join(mountpoint, 'My Drive'))
  if not force_remount and already_mounted:
    print(
        'Drive already mounted at {mnt}; to attempt to forcibly remount, '
        'call drive.mount("{mnt}", force_remount=True).'.format(mnt=mountpoint))
    return

  env = _env()
  home = env.home
  root_dir = env.root_dir
  inet_family = env.inet_family
  dev = env.dev
  path = env.path
  config_dir = env.config_dir

  try:
    _os.makedirs(config_dir)
  except OSError:
    if not _os.path.isdir(config_dir):
      raise ValueError('{} must be a directory if present'.format(config_dir))

  # Launch an intermediate bash to manage DriveFS' I/O (b/141747058#comment6).
  prompt = u'root@{}-{}: '.format(_socket.gethostname(), _uuid.uuid4().hex)
  logfile = None
  if mount._DEBUG:  # pylint:disable=protected-access
    logfile = _sys.stdout
  d = _popen_spawn.PopenSpawn(
      '/usr/bin/setsid /bin/bash --noediting -i',  # Need -i to get prompt echo.
      timeout=120,
      maxread=int(1e6),
      encoding='utf-8',
      logfile=logfile,
      env={
          'HOME': home,
          'FUSE_DEV_NAME': dev,
          'PATH': path
      })
  d.sendline('export PS1="{}"'.format(prompt))
  d.expect(prompt)  # The new prompt.
  drive_dir = _os.path.join(root_dir, 'opt/google/drive')
  # Robustify to previously-running copies of drive. Don't only [pkill -9]
  # because that leaves enough cruft behind in the mount table that future
  # operations fail with "Transport endpoint is not connected".
  d.sendline('umount -f {mnt} || umount {mnt}; pkill -9 -x drive'.format(
      mnt=mountpoint))
  # Wait for above to be received, using the next prompt.
  d.expect(prompt)
  d.sendline('pkill -9 -f {d}/directoryprefetcher_binary'.format(d=drive_dir))
  d.expect(prompt)
  # Only check the mountpoint after potentially unmounting/pkill'ing above.
  try:
    if _os.path.islink(mountpoint):
      raise ValueError('Mountpoint must not be a symlink')
    if _os.path.isdir(mountpoint) and _os.listdir(mountpoint):
      raise ValueError('Mountpoint must not already contain files')
    if not _os.path.isdir(mountpoint) and _os.path.exists(mountpoint):
      raise ValueError('Mountpoint must either be a directory or not exist')
    normed = _os.path.normpath(mountpoint)
    if '/' in normed and not _os.path.exists(_os.path.dirname(normed)):
      raise ValueError('Mountpoint must be in a directory that exists')
  except:
    d.kill(_signal.SIGKILL)
    raise

  # Watch for success.
  success = u'google.colab.drive MOUNTED'
  success_watcher = (
      '( while `sleep 0.5`; do if [[ -d "{m}" && "$(ls -A {m})" != "" ]]; '
      'then echo "{s}"; break; fi; done ) &').format(
          m=mountpoint, s=success)
  d.sendline(success_watcher)
  d.expect(prompt)

  oauth_prompt = u'(Go to this URL in a browser: https://.*)$'
  oauth_failed = u'Authorization failed'
  domain_disabled_drivefs = u'The domain policy has disabled Drive File Stream'
  problem_and_stopped = (
      u'Drive File Stream encountered a problem and has stopped')
  drive_exited = u'drive EXITED'
  metadata_auth_arg = (
      '--metadata_server_auth_uri={metadata_server}/computeMetadata/v1 '.format(
          metadata_server=_os.environ['TBE_CREDS_ADDR'])
      if use_metadata_server else '')

  # Create a pipe for sending the oauth code to a backgrounded drive binary.
  # (popen -> no pty -> no bash job control -> can't background post-launch).
  fifo_dir = _tempfile.mkdtemp()
  fifo = _os.path.join(fifo_dir, 'drive.fifo')
  _os.mkfifo(fifo)
  # cat is needed below since the FIFO isn't opened for writing yet.
  d.sendline((
      'cat {fifo} | head -1 | ( {d}/drive '
      '--features=' + ','.join([
          'fuse_max_background:1000',
          # TODO(b/168928162): Uncomment when b/176106731's fix rolls out.
          # 'max_bytes_per_fetch_content_request:134217728',
          'max_read_qps:1000',
          'max_write_qps:1000',
          'max_operation_batch_size:15',
          'max_parallel_push_task_instances:10',
          'opendir_timeout_ms:{timeout_ms}',
          'shortcut_support:true',
          'virtual_folders_omit_spaces:true',
      ]) + ' '
      '--inet_family=' + inet_family + ' ' + metadata_auth_arg +
      '--preferences=trusted_root_certs_file_path:'
      '{d}/roots.pem,mount_point_path:{mnt} --console_auth 2>&1 '
      '| grep --line-buffered -E "{oauth_prompt}|{problem_and_stopped}|{oauth_failed}|{domain_disabled_drivefs}"; '
      'echo "{drive_exited}"; ) &').format(
          d=drive_dir,
          timeout_ms=timeout_ms,
          mnt=mountpoint,
          fifo=fifo,
          oauth_failed=oauth_failed,
          domain_disabled_drivefs=domain_disabled_drivefs,
          oauth_prompt=oauth_prompt,
          problem_and_stopped=problem_and_stopped,
          drive_exited=drive_exited))
  d.expect(prompt)

  # LINT.IfChange(drivetimedout)
  timeout_pattern = 'QueryManager timed out'
  # LINT.ThenChange()
  dfs_log = _os.path.join(_logs_dir(), 'drive_fs.txt')

  # TODO(b/147296819): Delete this line.
  get_code = input if _sys.version_info[0] == 3 else raw_input  # pylint: disable=undefined-variable

  wrote_to_fifo = False
  while True:
    case = d.expect([
        success,
        prompt,
        oauth_prompt,
        problem_and_stopped,
        drive_exited,
        oauth_failed,
        domain_disabled_drivefs,
    ])
    if case == 0:
      break
    elif (case == 1 or case == 3 or case == 4):
      # Prompt appearing here means something went wrong with the drive binary.
      d.kill(_signal.SIGKILL)
      extra_reason = ''
      if 0 == _subprocess.call(
          'grep -q "{}" "{}"'.format(timeout_pattern, dfs_log), shell=True):
        extra_reason = (
            ': timeout during initial read of root folder; for more info: '
            'https://research.google.com/colaboratory/faq.html#drive-timeout')
      raise ValueError('mount failed' + extra_reason)
    elif case == 2:
      # Not already authorized, so do the authorization dance.
      auth_prompt = d.match.group(1) + '\nEnter your authorization code:\n'
      with _output.use_tags('dfs-auth-dance'):
        with open(fifo, 'w') as fifo_file:
          fifo_file.write(get_code(auth_prompt) + '\n')
      wrote_to_fifo = True
    elif case == 5:
      raise ValueError('mount failed: invalid oauth code')
    elif case == 6:
      raise ValueError(
          str(domain_disabled_drivefs) +
          ': https://support.google.com/a/answer/7496409')
  if not wrote_to_fifo:
    with open(fifo, 'w') as fifo_file:
      fifo_file.write('ignored\n')
  filtered_logfile = _timeouts_path()
  d.sendline('fuser -kw "{f}" ; rm -rf "{f}"'.format(f=filtered_logfile))
  d.expect(prompt)
  filter_script = _os.path.join(drive_dir, 'drive-filter.py')
  filter_cmd = (
      """nohup bash -c 'tail -n +0 -F "{}" | """
      """python3 {} > "{}" ' < /dev/null > /dev/null 2>&1 &""").format(
          dfs_log, filter_script, filtered_logfile)
  d.sendline(filter_cmd)
  d.expect(prompt)
  if 'ENABLE_DIRECTORYPREFETCHER' in _os.environ:
    d.sendline(
        """nohup bash -c '{d}/directoryprefetcher_binary -mountpoint={mnt}' """
        """>> {log} 2>&1 &""".format(
            d=drive_dir,
            mnt=mountpoint,
            log=_os.path.join(_logs_dir(), 'dpb.txt')))
    d.expect(prompt)
  d.sendline('disown -a')
  d.expect(prompt)
  d.sendline('exit')
  assert d.wait() == 0
  _output.clear(wait=True, output_tags='dfs-auth-dance')
  print('Mounted at {}'.format(mountpoint))

mount._DEBUG = False  # pylint:disable=protected-access
<commit_msg>Internal update.<commit_after># Copyright 2018 Google Inc.
#
# 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.
"""Colab-specific Google Drive integration."""

from __future__ import absolute_import as _
from __future__ import division as _
from __future__ import print_function as _

import collections as _collections
import os as _os
import signal as _signal
import socket as _socket
import subprocess as _subprocess
import sys as _sys
import tempfile as _tempfile
import uuid as _uuid

import pexpect.popen_spawn as _popen_spawn

from google.colab import output as _output

__all__ = ['flush_and_unmount', 'mount']

_Environment = _collections.namedtuple(
    '_Environment',
    ('home', 'root_dir', 'inet_family', 'dev', 'path', 'config_dir'))

def _env():
  """Create and return an _Environment to use."""
  home = _os.environ['HOME']
  root_dir = _os.path.realpath(
      _os.path.join(_os.environ['CLOUDSDK_CONFIG'], '../..'))
  inet_family = 'IPV4_ONLY'
  dev = '/dev/fuse'
  path = '/usr/local/bin:/usr/local/sbin:/usr/bin:/usr/sbin:/bin:/sbin:.'
  if len(root_dir) > 1 and not root_dir.startswith('/usr/local/google/'):
    home = _os.path.join(root_dir, home)
    inet_family = 'IPV6_ONLY'
    fum = _os.environ['HOME'].split('mount')[0] + '/mount/alloc/fusermount'
    if 'BORG_ALLOC_DIR' in _os.environ:
      # For support in guitar cluster invocations.
      fum = _os.environ['BORG_ALLOC_DIR'] + '/alloc/fusermount'
    dev = fum + '/dev/fuse'
    path = path + ':' + fum + '/bin'
  config_dir = _os.path.join(home, '.config', 'Google')
  return _Environment(
      home=home,
      root_dir=root_dir,
      inet_family=inet_family,
      dev=dev,
      path=path,
      config_dir=config_dir)

def _logs_dir():
  return _os.path.join(_env().config_dir, 'DriveFS/Logs/')

def _timeouts_path():
  return _os.path.join(_logs_dir(), 'timeouts.txt')

def flush_and_unmount(timeout_ms=24 * 60 * 60 * 1000):
  """Unmount Google Drive and flush any outstanding writes to it."""
  env = _env()
  if b'type fuse.drive' not in _subprocess.check_output(['/bin/mount']):
    print('Drive not mounted, so nothing to flush and unmount.')
    return
  drive_bin = _os.path.join(env.root_dir, 'opt/google/drive/drive')
  p = _subprocess.Popen([
      drive_bin, '--push_changes_and_quit', '--single_process',
      '--timeout_sec={}'.format(int(timeout_ms / 1000))
  ],
                        stdout=_subprocess.PIPE,
                        stderr=_subprocess.PIPE)
  out, err = p.communicate()
  if mount._DEBUG:  # pylint:disable=protected-access
    print('flush_and_unmount: out: {}\nerr: {}'.format(out, err))
  if p.returncode:
    raise ValueError('flush_and_unmount failed')

def mount(mountpoint,
          force_remount=False,
          timeout_ms=60000,
          use_metadata_server=False):
  """Mount your Google Drive at the specified mountpoint path."""

  if ' ' in mountpoint:
    raise ValueError('Mountpoint must not contain a space.')

  mountpoint = _os.path.expanduser(mountpoint)
  # If we've already mounted drive at the specified mountpoint, exit now.
  already_mounted = _os.path.isdir(_os.path.join(mountpoint, 'My Drive'))
  if not force_remount and already_mounted:
    print(
        'Drive already mounted at {mnt}; to attempt to forcibly remount, '
        'call drive.mount("{mnt}", force_remount=True).'.format(mnt=mountpoint))
    return

  env = _env()
  home = env.home
  root_dir = env.root_dir
  inet_family = env.inet_family
  dev = env.dev
  path = env.path
  config_dir = env.config_dir

  try:
    _os.makedirs(config_dir)
  except OSError:
    if not _os.path.isdir(config_dir):
      raise ValueError('{} must be a directory if present'.format(config_dir))

  # Launch an intermediate bash to manage DriveFS' I/O (b/141747058#comment6).
  prompt = u'root@{}-{}: '.format(_socket.gethostname(), _uuid.uuid4().hex)
  logfile = None
  if mount._DEBUG:  # pylint:disable=protected-access
    logfile = _sys.stdout
  d = _popen_spawn.PopenSpawn(
      '/usr/bin/setsid /bin/bash --noediting -i',  # Need -i to get prompt echo.
      timeout=120,
      maxread=int(1e6),
      encoding='utf-8',
      logfile=logfile,
      env={
          'HOME': home,
          'FUSE_DEV_NAME': dev,
          'PATH': path
      })
  d.sendline('unset HISTFILE; export PS1="{}"'.format(prompt))
  d.expect(prompt)  # The new prompt.
  drive_dir = _os.path.join(root_dir, 'opt/google/drive')
  # Robustify to previously-running copies of drive. Don't only [pkill -9]
  # because that leaves enough cruft behind in the mount table that future
  # operations fail with "Transport endpoint is not connected".
  d.sendline('umount -f {mnt} || umount {mnt}; pkill -9 -x drive'.format(
      mnt=mountpoint))
  # Wait for above to be received, using the next prompt.
  d.expect(prompt)
  d.sendline('pkill -9 -f {d}/directoryprefetcher_binary'.format(d=drive_dir))
  d.expect(prompt)
  # Only check the mountpoint after potentially unmounting/pkill'ing above.
  try:
    if _os.path.islink(mountpoint):
      raise ValueError('Mountpoint must not be a symlink')
    if _os.path.isdir(mountpoint) and _os.listdir(mountpoint):
      raise ValueError('Mountpoint must not already contain files')
    if not _os.path.isdir(mountpoint) and _os.path.exists(mountpoint):
      raise ValueError('Mountpoint must either be a directory or not exist')
    normed = _os.path.normpath(mountpoint)
    if '/' in normed and not _os.path.exists(_os.path.dirname(normed)):
      raise ValueError('Mountpoint must be in a directory that exists')
  except:
    d.kill(_signal.SIGKILL)
    raise

  # Watch for success.
  success = u'google.colab.drive MOUNTED'
  success_watcher = (
      '( while `sleep 0.5`; do if [[ -d "{m}" && "$(ls -A {m})" != "" ]]; '
      'then echo "{s}"; break; fi; done ) &').format(
          m=mountpoint, s=success)
  d.sendline(success_watcher)
  d.expect(prompt)

  oauth_prompt = u'(Go to this URL in a browser: https://.*)$'
  oauth_failed = u'Authorization failed'
  domain_disabled_drivefs = u'The domain policy has disabled Drive File Stream'
  problem_and_stopped = (
      u'Drive File Stream encountered a problem and has stopped')
  drive_exited = u'drive EXITED'
  metadata_auth_arg = (
      '--metadata_server_auth_uri={metadata_server}/computeMetadata/v1 '.format(
          metadata_server=_os.environ['TBE_CREDS_ADDR'])
      if use_metadata_server else '')

  # Create a pipe for sending the oauth code to a backgrounded drive binary.
  # (popen -> no pty -> no bash job control -> can't background post-launch).
  fifo_dir = _tempfile.mkdtemp()
  fifo = _os.path.join(fifo_dir, 'drive.fifo')
  _os.mkfifo(fifo)
  # cat is needed below since the FIFO isn't opened for writing yet.
  d.sendline((
      'cat {fifo} | head -1 | ( {d}/drive '
      '--features=' + ','.join([
          'fuse_max_background:1000',
          # TODO(b/168928162): Uncomment when b/176106731's fix rolls out.
          # 'max_bytes_per_fetch_content_request:134217728',
          'max_read_qps:1000',
          'max_write_qps:1000',
          'max_operation_batch_size:15',
          'max_parallel_push_task_instances:10',
          'opendir_timeout_ms:{timeout_ms}',
          'shortcut_support:true',
          'virtual_folders_omit_spaces:true',
      ]) + ' '
      '--inet_family=' + inet_family + ' ' + metadata_auth_arg +
      '--preferences=trusted_root_certs_file_path:'
      '{d}/roots.pem,mount_point_path:{mnt} --console_auth 2>&1 '
      '| grep --line-buffered -E "{oauth_prompt}|{problem_and_stopped}|{oauth_failed}|{domain_disabled_drivefs}"; '
      'echo "{drive_exited}"; ) &').format(
          d=drive_dir,
          timeout_ms=timeout_ms,
          mnt=mountpoint,
          fifo=fifo,
          oauth_failed=oauth_failed,
          domain_disabled_drivefs=domain_disabled_drivefs,
          oauth_prompt=oauth_prompt,
          problem_and_stopped=problem_and_stopped,
          drive_exited=drive_exited))
  d.expect(prompt)

  # LINT.IfChange(drivetimedout)
  timeout_pattern = 'QueryManager timed out'
  # LINT.ThenChange()
  dfs_log = _os.path.join(_logs_dir(), 'drive_fs.txt')

  # TODO(b/147296819): Delete this line.
  get_code = input if _sys.version_info[0] == 3 else raw_input  # pylint: disable=undefined-variable

  wrote_to_fifo = False
  while True:
    case = d.expect([
        success,
        prompt,
        oauth_prompt,
        problem_and_stopped,
        drive_exited,
        oauth_failed,
        domain_disabled_drivefs,
    ])
    if case == 0:
      break
    elif (case == 1 or case == 3 or case == 4):
      # Prompt appearing here means something went wrong with the drive binary.
      d.kill(_signal.SIGKILL)
      extra_reason = ''
      if 0 == _subprocess.call(
          'grep -q "{}" "{}"'.format(timeout_pattern, dfs_log), shell=True):
        extra_reason = (
            ': timeout during initial read of root folder; for more info: '
            'https://research.google.com/colaboratory/faq.html#drive-timeout')
      raise ValueError('mount failed' + extra_reason)
    elif case == 2:
      # Not already authorized, so do the authorization dance.
      auth_prompt = d.match.group(1) + '\nEnter your authorization code:\n'
      with _output.use_tags('dfs-auth-dance'):
        with open(fifo, 'w') as fifo_file:
          fifo_file.write(get_code(auth_prompt) + '\n')
      wrote_to_fifo = True
    elif case == 5:
      raise ValueError('mount failed: invalid oauth code')
    elif case == 6:
      raise ValueError(
          str(domain_disabled_drivefs) +
          ': https://support.google.com/a/answer/7496409')
  if not wrote_to_fifo:
    with open(fifo, 'w') as fifo_file:
      fifo_file.write('ignored\n')
  filtered_logfile = _timeouts_path()
  d.sendline('fuser -kw "{f}" ; rm -rf "{f}"'.format(f=filtered_logfile))
  d.expect(prompt)
  filter_script = _os.path.join(drive_dir, 'drive-filter.py')
  filter_cmd = (
      """nohup bash -c 'tail -n +0 -F "{}" | """
      """python3 {} > "{}" ' < /dev/null > /dev/null 2>&1 &""").format(
          dfs_log, filter_script, filtered_logfile)
  d.sendline(filter_cmd)
  d.expect(prompt)
  if 'ENABLE_DIRECTORYPREFETCHER' in _os.environ:
    d.sendline(
        """nohup bash -c '{d}/directoryprefetcher_binary -mountpoint={mnt}' """
        """>> {log} 2>&1 &""".format(
            d=drive_dir,
            mnt=mountpoint,
            log=_os.path.join(_logs_dir(), 'dpb.txt')))
    d.expect(prompt)
  d.sendline('disown -a')
  d.expect(prompt)
  d.sendline('exit')
  assert d.wait() == 0
  _output.clear(wait=True, output_tags='dfs-auth-dance')
  print('Mounted at {}'.format(mountpoint))

mount._DEBUG = False  # pylint:disable=protected-access
<|endoftext|>


<commit_before># Note: Work in progress

import os
import re
import time
import codecs
from StringIO import StringIO

import Version
from Code import CCodeWriter
from Cython import Utils

# need one-characters subsitutions (for now) so offsets aren't off
special_chars = [(u'<', u'\xF0', u'<'),
                 (u'>', u'\xF1', u'>'),
                 (u'&', u'\xF2', u'&')]

class AnnotationCCodeWriter(CCodeWriter):

    def __init__(self, create_from=None, buffer=None, copy_formatting=True):
        CCodeWriter.__init__(self, create_from, buffer, copy_formatting=True)
        if create_from is None:
            self.annotation_buffer = StringIO()
            self.annotations = []
            self.last_pos = None
            self.code = {}
        else:
            # When creating an insertion point, keep references to the same database
            self.annotation_buffer = create_from.annotation_buffer
            self.annotations = create_from.annotations
            self.code = create_from.code
            self.last_pos = create_from.last_pos

    def create_new(self, create_from, buffer, copy_formatting):
        return AnnotationCCodeWriter(create_from, buffer, copy_formatting)

    def write(self, s):
        CCodeWriter.write(self, s)
        self.annotation_buffer.write(s)

    def mark_pos(self, pos):
#        if pos is not None:
#            CCodeWriter.mark_pos(self, pos)
#        return
        if self.last_pos:
            code = self.code.get(self.last_pos[1], "")
            self.code[self.last_pos[1]] = code + self.annotation_buffer.getvalue()
        self.annotation_buffer = StringIO()
        self.last_pos = pos

    def annotate(self, pos, item):
        self.annotations.append((pos, item))

    def save_annotation(self, source_filename, target_filename):
        self.mark_pos(None)
        f = Utils.open_source_file(source_filename)
        lines = f.readlines()
        for k in range(len(lines)):
            line = lines[k]
            for c, cc, html in special_chars:
                line = line.replace(c, cc)
            lines[k] = line
        f.close()
        all = []
        for pos, item in self.annotations:
            if pos[0] == source_filename:
                start = item.start()
                size, end = item.end()
                if size:
                    all.append((pos, start))
                    all.append(((source_filename, pos[1], pos[2]+size), end))
                else:
                    all.append((pos, start+end))

        all.sort()
        all.reverse()
        for pos, item in all:
            _, line_no, col = pos
            line_no -= 1
            col += 1
            line = lines[line_no]
            lines[line_no] = line[:col] + item + line[col:]

        html_filename = os.path.splitext(target_filename)[0] + ".html"
        f = codecs.open(html_filename, "w", encoding="UTF-8")
        f.write(u'\n')
        f.write(u"""

<meta http-equiv="Content-Type" content="text/html; charset=utf-8" />

        """)
        f.write(u'\n')
        f.write(u'
Generated by Cython %s on %s\n' % (Version.version, time.asctime()))
        c_file = Utils.encode_filename(os.path.basename(target_filename))
        f.write(u'
Raw output: %s\n' % (c_file, c_file))
        k = 0

        py_c_api = re.compile(u'(Py[A-Z][a-z]+_[A-Z][a-z][A-Za-z_]+)')
        pyx_api = re.compile(u'(__Pyx[A-Za-z_]+)\(')
        py_marco_api = re.compile(u'(Py[A-Za-z]*_[A-Z][A-Z_]+)')
        error_goto = re.compile(ur'((; *if .*)? \{__pyx_filename = .*goto __pyx_L\w+;\})')

        for line in lines:

            k += 1
            try:
                code = self.code[k]
            except KeyError:
                code = ''

            code, c_api_calls = py_c_api.subn(ur"\1", code)
            code, pyx_api_calls = pyx_api.subn(ur"\1(", code)
            code, macro_api_calls = py_marco_api.subn(ur"\1", code)
            code, error_goto_calls = error_goto.subn(ur"\1", code)

            code = code.replace(u";", u";")

            color = u"FFFF%02x" % int(255/(1+(5*c_api_calls+2*pyx_api_calls+macro_api_calls)/10.0))
            f.write(u"
" % (color, k))

            f.write(u" %d: " % k)
            for c, cc, html in special_chars:
                line = line.replace(cc, html)
            f.write(line.rstrip())

            f.write(u'
\n')
            f.write(u"
%s
" % (k, color, code))
        f.write(u'\n')
        f.close()

# TODO: make this cleaner
def escape(raw_string):
    raw_string = raw_string.replace(u"\'", ur"’")
    raw_string = raw_string.replace(u'\"', ur'"')
    raw_string = raw_string.replace(u'\n', ur'
\n')
    raw_string = raw_string.replace(u'\t', ur'\t')
    return raw_string

class AnnotationItem:

    def __init__(self, style, text, tag="", size=0):
        self.style = style
        self.text = text
        self.tag = tag
        self.size = size

    def start(self):
        return u"%s" % (self.style, self.text, self.tag)

    def end(self):
        return self.size, u""
<commit_msg>fix Cython code context comments in C code when annotation is enables<commit_after># Note: Work in progress

import os
import re
import time
import codecs
from StringIO import StringIO

import Version
from Code import CCodeWriter
from Cython import Utils

# need one-characters subsitutions (for now) so offsets aren't off
special_chars = [(u'<', u'\xF0', u'<'),
                 (u'>', u'\xF1', u'>'),
                 (u'&', u'\xF2', u'&')]

class AnnotationCCodeWriter(CCodeWriter):

    def __init__(self, create_from=None, buffer=None, copy_formatting=True):
        CCodeWriter.__init__(self, create_from, buffer, copy_formatting=True)
        if create_from is None:
            self.annotation_buffer = StringIO()
            self.annotations = []
            self.last_pos = None
            self.code = {}
        else:
            # When creating an insertion point, keep references to the same database
            self.annotation_buffer = create_from.annotation_buffer
            self.annotations = create_from.annotations
            self.code = create_from.code
            self.last_pos = create_from.last_pos

    def create_new(self, create_from, buffer, copy_formatting):
        return AnnotationCCodeWriter(create_from, buffer, copy_formatting)

    def write(self, s):
        CCodeWriter.write(self, s)
        self.annotation_buffer.write(s)

    def mark_pos(self, pos):
        if pos is not None:
            CCodeWriter.mark_pos(self, pos)
        if self.last_pos:
            code = self.code.get(self.last_pos[1], "")
            self.code[self.last_pos[1]] = code + self.annotation_buffer.getvalue()
        self.annotation_buffer = StringIO()
        self.last_pos = pos

    def annotate(self, pos, item):
        self.annotations.append((pos, item))

    def save_annotation(self, source_filename, target_filename):
        self.mark_pos(None)
        f = Utils.open_source_file(source_filename)
        lines = f.readlines()
        for k in range(len(lines)):
            line = lines[k]
            for c, cc, html in special_chars:
                line = line.replace(c, cc)
            lines[k] = line
        f.close()
        all = []
        for pos, item in self.annotations:
            if pos[0] == source_filename:
                start = item.start()
                size, end = item.end()
                if size:
                    all.append((pos, start))
                    all.append(((source_filename, pos[1], pos[2]+size), end))
                else:
                    all.append((pos, start+end))

        all.sort()
        all.reverse()
        for pos, item in all:
            _, line_no, col = pos
            line_no -= 1
            col += 1
            line = lines[line_no]
            lines[line_no] = line[:col] + item + line[col:]

        html_filename = os.path.splitext(target_filename)[0] + ".html"
        f = codecs.open(html_filename, "w", encoding="UTF-8")
        f.write(u'\n')
        f.write(u"""

<meta http-equiv="Content-Type" content="text/html; charset=utf-8" />

        """)
        f.write(u'\n')
        f.write(u'
Generated by Cython %s on %s\n' % (Version.version, time.asctime()))
        c_file = Utils.encode_filename(os.path.basename(target_filename))
        f.write(u'
Raw output: %s\n' % (c_file, c_file))
        k = 0

        py_c_api = re.compile(u'(Py[A-Z][a-z]+_[A-Z][a-z][A-Za-z_]+)')
        pyx_api = re.compile(u'(__Pyx[A-Za-z_]+)\(')
        py_marco_api = re.compile(u'(Py[A-Za-z]*_[A-Z][A-Z_]+)')
        error_goto = re.compile(ur'((; *if .*)? \{__pyx_filename = .*goto __pyx_L\w+;\})')

        for line in lines:

            k += 1
            try:
                code = self.code[k]
            except KeyError:
                code = ''

            code, c_api_calls = py_c_api.subn(ur"\1", code)
            code, pyx_api_calls = pyx_api.subn(ur"\1(", code)
            code, macro_api_calls = py_marco_api.subn(ur"\1", code)
            code, error_goto_calls = error_goto.subn(ur"\1", code)

            code = code.replace(u";", u";")

            color = u"FFFF%02x" % int(255/(1+(5*c_api_calls+2*pyx_api_calls+macro_api_calls)/10.0))
            f.write(u"
" % (color, k))

            f.write(u" %d: " % k)
            for c, cc, html in special_chars:
                line = line.replace(cc, html)
            f.write(line.rstrip())

            f.write(u'
\n')
            f.write(u"
%s
" % (k, color, code))
        f.write(u'\n')
        f.close()

# TODO: make this cleaner
def escape(raw_string):
    raw_string = raw_string.replace(u"\'", ur"’")
    raw_string = raw_string.replace(u'\"', ur'"')
    raw_string = raw_string.replace(u'\n', ur'
\n')
    raw_string = raw_string.replace(u'\t', ur'\t')
    return raw_string

class AnnotationItem:

    def __init__(self, style, text, tag="", size=0):
        self.style = style
        self.text = text
        self.tag = tag
        self.size = size

    def start(self):
        return u"%s" % (self.style, self.text, self.tag)

    def end(self):
        return self.size, u""
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
:Author Patrik Valkovic
:Created 22.08.2017 11:37
:Licence GNUv3
Part of grammpy-transforms

"""

from grammpy import *

def remove_unit_rules(grammar: Grammar, transform_grammar=False) -> Grammar:
    raise NotImplementedError()<commit_msg>Add ReducedUnitRule class<commit_after>#!/usr/bin/env python
"""
:Author Patrik Valkovic
:Created 22.08.2017 11:37
:Licence GNUv3
Part of grammpy-transforms

"""

from typing import List
from grammpy import *

class ReducedUnitRule(Rule):
    by_rules = []  # type: List[Rule]
    end_rule = None  # type: Rule

def remove_unit_rules(grammar: Grammar, transform_grammar=False) -> Grammar:
    raise NotImplementedError()<|endoftext|>


<commit_before>from django.core.urlresolvers import reverse
from django.db import models
from django.db.models.signals import pre_save
from django.dispatch import receiver
from django.template.defaultfilters import slugify, stringformat
import markdown
import re
from bs4 import BeautifulSoup

class Tag(models.Model):
    name = models.CharField(max_length=50, unique=True)
    slug = models.SlugField(max_length=50, unique=True)

    def __str__(self):
        return self.name

class Post(models.Model):
    title = models.CharField(max_length=200)
    date = models.DateTimeField()
    content = models.TextField()
    content_html = models.TextField(blank=True)
    tags = models.ManyToManyField(Tag, related_name='posts', blank=True)
    slug = models.SlugField(max_length=200, unique=True, blank=True)

    def get_absolute_url(self):
        return reverse('blog:post', kwargs={
            'year': self.date.year,
            'month': stringformat(self.date.month, '02d'),
            'day': stringformat(self.date.day, '02d'),
            'slug': self.slug
        })

    def __str__(self):
        return self.title

@receiver(pre_save, sender=Post)
def pre_save_post(**kwargs):
    content_html = markdown.markdown(kwargs['instance'].content,
                                     extensions=['codehilite'])
    soup = BeautifulSoup(content_html)
    for tag in soup.find_all(re.compile(r'h\d')):
        if tag.parent is soup:
            tag.name = 'h%d' % (int(tag.name[1]) + 1)
    kwargs['instance'].content_html = soup.prettify()
    kwargs['instance'].slug = slugify(kwargs['instance'].title)
<commit_msg>Don't pretty print markdown output<commit_after>from django.core.urlresolvers import reverse
from django.db import models
from django.db.models.signals import pre_save
from django.dispatch import receiver
from django.template.defaultfilters import slugify, stringformat
import markdown
import re
from bs4 import BeautifulSoup

class Tag(models.Model):
    name = models.CharField(max_length=50, unique=True)
    slug = models.SlugField(max_length=50, unique=True)

    def __str__(self):
        return self.name

class Post(models.Model):
    title = models.CharField(max_length=200)
    date = models.DateTimeField()
    content = models.TextField()
    content_html = models.TextField(blank=True)
    tags = models.ManyToManyField(Tag, related_name='posts', blank=True)
    slug = models.SlugField(max_length=200, unique=True, blank=True)

    def get_absolute_url(self):
        return reverse('blog:post', kwargs={
            'year': self.date.year,
            'month': stringformat(self.date.month, '02d'),
            'day': stringformat(self.date.day, '02d'),
            'slug': self.slug
        })

    def __str__(self):
        return self.title

@receiver(pre_save, sender=Post)
def pre_save_post(**kwargs):
    content_html = markdown.markdown(kwargs['instance'].content,
                                     extensions=['codehilite'])
    soup = BeautifulSoup(content_html)
    for tag in soup.find_all(re.compile(r'h\d')):
        if tag.parent is soup:
            tag.name = 'h%d' % (int(tag.name[1]) + 1)
    kwargs['instance'].content_html = str(soup)
    kwargs['instance'].slug = slugify(kwargs['instance'].title)
<|endoftext|>


<commit_before># cython: infer_types=True, language_level=3, py2_import=True
#
#   Cython Scanner
#

from __future__ import absolute_import

import cython
cython.declare(EncodedString=object, make_lexicon=object, lexicon=object,
               any_string_prefix=unicode, IDENT=unicode,
               print_function=object, error=object, warning=object,
               os=object, platform=object)

import os
import platform

from .. import Utils
from ..Plex.Scanners import Scanner
from ..Plex.Errors import UnrecognizedInput
from .Errors import error, warning
from .Lexicon import any_string_prefix, make_lexicon, IDENT
from .Future import print_function

from .StringEncoding import EncodedString

debug_scanner = 0
trace_scanner = 0
scanner_debug_flags = 0
scanner_dump_file = None

lexicon = None

def get_lexicon():
    global lexicon
    if not lexicon:
        lexicon = make_lexicon()
    return lexicon

#------------------------------------------------------------------

py_reserved_words = [
    "global", "nonlocal", "def", "class", "print", "del", "pass", "break",
    "continue", "return", "raise", "import", "exec", "try",
    "except", "finally", "while", "if", "elif", "else", "for",
    "in", "assert", "and", "or", "not", "is", "in", "lambda",
    "from", "yield", "with", "nonlocal",
]

pyx_reserved_words = py_reserved_words + [
    "include", "ctypedef", "cdef", "cpdef",
    "cimport", "DEF", "IF", "ELIF", "ELSE"
]

class Method(object):

    def __init__(self, name):
        self.name = name
        self.__name__ = name  # for Plex tracing

    def __call__(self, stream, text):
        return getattr(stream, self.name)(text)

#------------------------------------------------------------------

class CompileTimeScope(object):

    def __init__(self, outer = None):
        self.entries = {}
        self.outer = outer

    def declare(self, name, value):
        self.entries[name] = value

    def update(self, other):
        self.entries.update(other)

    def lookup_here(self, name):
        return self.entries[name]

    def __contains__(self, name):
        return name in self.entries

    def lookup(self, name):
        try:
            return self.lookup_here(name)
        except KeyError:
            outer = self.outer
            if outer:
                return outer.lookup(name)
            else:
                raise

def initial_compile_time_env():
    benv = CompileTimeScope()
    names = ('UNAME_SYSNAME', 'UNAME_NODENAME', 'UNAME_RELEASE',
        'UNAME_VERSION', 'UNAME_MACHINE')
    for name, value in zip(names, platform.uname()):
        benv.declare(name, value)
    try:
        import __builtin__ as builtins
    except ImportError:
        import builtins

    names = ('False', 'True',
             'abs', 'all', 'any', 'ascii', 'bin', 'bool', 'bytearray', 'bytes',
             'chr', 'cmp', 'complex', 'dict', 'divmod', 'enumerate', 'filter',
             'float', 'format', 'frozenset', 'hash', 'hex', 'int', 'len',
             'list', 'long', 'map', 'max', 'min', 'oct', 'ord', 'pow', 'range',
             'repr', 'reversed', 'round', 'set', 'slice', 'sorted', 'str',
             'sum', 'tuple', 'xrange', 'zip')

    for name in names:
        try:
            benv.declare(name, getattr(builtins, name))
        except AttributeError:
            # ignore, likely Py3
            pass
    denv = CompileTimeScope(benv)
    return denv

#------------------------------------------------------------------

class SourceDescriptor(object):
    """
    A SourceDescriptor should be considered immutable.
    """
    _file_type = 'pyx'

    _escaped_description = None
    _cmp_name = ''
    def __str__(self):
        assert False # To catch all places where a descriptor is used directly as a filename

    def set_file_type_from_name(self, filename):
        name, ext = os.path.splitext(filename)
        self._file_type = ext in ('.pyx', '.pxd', '.py') and ext[1:] or 'pyx'

    def is_cython_file(self):
        return self._file_type in ('pyx', 'pxd')

    def is_python_file(self):
        return self._file_type == 'py'

    def get_escaped_description(self):
        if self._escaped_description is None:
            self._escaped_description = \
                self.get_description().encode('ASCII', 'replace').decode("ASCII")
        return self._escaped_description

    def __gt__(self, other):
        # this is only used to provide some sort of order
        try:
            return self._cmp_name > other._cmp_name
        except AttributeError:
            return False

    def __lt__(self, other):
        # this is only used to provide some sort of order
        try:
            return self._cmp_name < other._cmp_name
        except AttributeError:
            return False

    def __le__(self, other):
        # this is only used to provide some sort of order
        try:
            return self._cmp_name <= other._cmp_name
        except AttributeError:
            return False

class FileSourceDescriptor(SourceDescriptor):
    """
    Represents a code source. A code source is a more generic abstraction
    for a "filename" (as sometimes the code doesn't come from a file).
    Instances of code sources are passed to Scanner.__init__ as the
    optional name argument and will be passed back when asking for
    the position()-tuple.
    """
    def __init__(self, filename, path_description=None):
        filename = Utils.decode_filename(filename)
        self.path_description = path_description or filename
        self.filename = filename
        self.set_file_type_from_name(filename)
        self._cmp_name = filename
        self._lines = {}

    def get_lines(self, encoding=None, error_handling=None):
        # we cache the lines only the second time this is called, in
        # order to save memory when they are only used once
        key = (encoding, error_handling)
        try:
            lines = self._lines[key]
            if lines is not None:
                return lines
        except KeyError:
            pass
        f = Utils.open_source_file(
            self.filename, encoding=encoding,
            error_handling=error_handling,
            # newline normalisation is costly before Py2.6
            require_normalised_newlines=False)
        try:
            lines = list(f)
        finally:
            f.close()
        if key in self._lines:
            self._lines[key] = lines
        else:
            # do not cache the first access, but remember that we
            # already read it once
            self._lines[key] = None
        return lines

    def get_description(self):
        try:
            return os.path.relpath(self.path_description)
        except ValueError:
            # path not under current directory => use complete file path
            return self.path_description

    def get_error_description(self):
        path = self.filename
        cwd = Utils.decode_filename(os.getcwd() + os.path.sep)
        if path.startswith(cwd):
            return path[len(cwd):]
        return path

    def get_filenametable_entry(self):
        return self.filename

    def __eq__(self, other):
        return isinstance(other, FileSourceDescriptor) and self.filename == other.filename

    def __hash__(self):
        return hash(self.filename)

    def __repr__(self):
        return "<FileSourceDescriptor:%s>" % self.filename

class StringSourceDescriptor(SourceDescriptor):
    """
    Instances of this class can be used instead of a filenames if the
    code originates from a string object.
    """
    filename = None

    def __init__(self, name, code):
        self.name = name
        #self.set_file_type_from_name(name)
        self.codelines = [x + "\n" for x in code.split("\n")]
        self._cmp_name = name

    def get_lines(self, encoding=None, error_handling=None):
        if not encoding:
            return self.codelines
        else:
            return [ line.encode(encoding, error_handling).decode(encoding)
                     for line in self.codelines ]

    def get_description(self):
        return self.name

    get_error_description = get_description

    def get_filenametable_entry(self):
        return "stringsource"

    def __hash__(self):
        return id(self)
        # Do not hash on the name, an identical string source should be the
        # same object (name is often defaulted in other places)
        # return hash(self.name)

    def __eq__(self, other):
        return isinstance(other, StringSourceDescriptor) and self.name == other.name

    def __repr__(self):
        return "<StringSourceDescriptor:%s>" % self.name

#------------------------------------------------------------------

class PyrexScanner(Scanner):
    #  context            Context  Compilation context
    #  included_files     [string] Files included with 'include' statement
    #  compile_time_env   dict     Environment for conditional compilation
    #  compile_time_eval  boolean  In a true conditional compilation context
    #  compile_time_expr  boolean  In a compile-time expression context

    def __init__(self, file, filename, parent_scanner=None,
                 scope=None, context=None, source_encoding=None, parse_comments=True, initial_pos=None):
        Scanner.__init__(self, get_lexicon(), file, filename, initial_pos)
        if parent_scanner:
            self.context = parent_scanner.context
            self.included_files = parent_scanner.included_files
            self.compile_time_env = parent_scanner.compile_time_env
            self.compile_time_eval = parent_scanner.compile_time_eval
            self.compile_time_expr = parent_scanner.compile_time_expr
        else:
            self.context = context
            self.included_files = scope.included_files
            self.compile_time_env = initial_compile_time_env()
            self.compile_time_eval = 1
            self.compile_time_expr = 0
            if getattr(context.options, 'compile_time_env', None):
                self.compile_time_env.update(context.options.compile_time_env)
        self.parse_comments = parse_comments
        self.source_encoding = source_encoding
        if filename.is_python_file():
            self.in_python_file = True
            self.keywords = set(py_reserved_words)
        else:
            self.in_python_file = False
            self.keywords = set(pyx_reserved_words)
        self.trace = trace_scanner
        self.indentation_stack = [0]
        self.indentation_char = None
        self.bracket_nesting_level = 0
        self.begin('INDENT')
        self.sy = ''
        self.next()

    def commentline(self, text):
        if self.parse_comments:
            self.produce('commentline', text)

    def current_level(self):
        return self.indentation_stack[-1]

    def open_bracket_action(self, text):
        self.bracket_nesting_level += 1
        return text

    def close_bracket_action(self, text):
        self.bracket_nesting_level -= 1
        return text

    def newline_action(self, text):
        if self.bracket_nesting_level == 0:
            self.begin('INDENT')
            self.produce('NEWLINE', '')

    string_states = {
        "'":   'SQ_STRING',
        '"':   'DQ_STRING',
        "'''": 'TSQ_STRING',
        '"""': 'TDQ_STRING'
    }

    def begin_string_action(self, text):
        while text[:1] in any_string_prefix:
            text = text[1:]
        self.begin(self.string_states[text])
        self.produce('BEGIN_STRING')

    def end_string_action(self, text):
        self.begin('')
        self.produce('END_STRING')

    def unclosed_string_action(self, text):
        self.end_string_action(text)
        self.error("Unclosed string literal")

    def indentation_action(self, text):
        self.begin('')
        # Indentation within brackets should be ignored.
        #if self.bracket_nesting_level > 0:
        #    return
        # Check that tabs and spaces are being used consistently.
        if text:
            c = text[0]
            #print "Scanner.indentation_action: indent with", repr(c) ###
            if self.indentation_char is None:
                self.indentation_char = c
                #print "Scanner.indentation_action: setting indent_char to", repr(c)
            else:
                if self.indentation_char != c:
                    self.error("Mixed use of tabs and spaces")
            if text.replace(c, "") != "":
                self.error("Mixed use of tabs and spaces")
        # Figure out how many indents/dedents to do
        current_level = self.current_level()
        new_level = len(text)
        #print "Changing indent level from", current_level, "to", new_level ###
        if new_level == current_level:
            return
        elif new_level > current_level:
            #print "...pushing level", new_level ###
            self.indentation_stack.append(new_level)
            self.produce('INDENT', '')
        else:
            while new_level < self.current_level():
                #print "...popping level", self.indentation_stack[-1] ###
                self.indentation_stack.pop()
                self.produce('DEDENT', '')
            #print "...current level now", self.current_level() ###
            if new_level != self.current_level():
                self.error("Inconsistent indentation")

    def eof_action(self, text):
        while len(self.indentation_stack) > 1:
            self.produce('DEDENT', '')
            self.indentation_stack.pop()
        self.produce('EOF', '')

    def next(self):
        try:
            sy, systring = self.read()
        except UnrecognizedInput:
            self.error("Unrecognized character")
            return  # just a marker, error() always raises
        if sy == IDENT:
            if systring in self.keywords:
                if systring == u'print' and print_function in self.context.future_directives:
                    self.keywords.discard('print')
                elif systring == u'exec' and self.context.language_level >= 3:
                    self.keywords.discard('exec')
                else:
                    sy = systring
            systring = self.context.intern_ustring(systring)
        self.sy = sy
        self.systring = systring
        if False: # debug_scanner:
            _, line, col = self.position()
            if not self.systring or self.sy == self.systring:
                t = self.sy
            else:
                t = "%s %s" % (self.sy, self.systring)
            print("--- %3d %2d %s" % (line, col, t))

    def peek(self):
        saved = self.sy, self.systring
        self.next()
        next = self.sy, self.systring
        self.unread(*next)
        self.sy, self.systring = saved
        return next

    def put_back(self, sy, systring):
        self.unread(self.sy, self.systring)
        self.sy = sy
        self.systring = systring

    def unread(self, token, value):
        # This method should be added to Plex
        self.queue.insert(0, (token, value))

    def error(self, message, pos=None, fatal=True):
        if pos is None:
            pos = self.position()
        if self.sy == 'INDENT':
            error(pos, "Possible inconsistent indentation")
        err = error(pos, message)
        if fatal: raise err

    def expect(self, what, message=None):
        if self.sy == what:
            self.next()
        else:
            self.expected(what, message)

    def expect_keyword(self, what, message=None):
        if self.sy == IDENT and self.systring == what:
            self.next()
        else:
            self.expected(what, message)

    def expected(self, what, message = None):
        if message:
            self.error(message)
        else:
            if self.sy == IDENT:
                found = self.systring
            else:
                found = self.sy
            self.error("Expected '%s', found '%s'" % (what, found))

    def expect_indent(self):
        self.expect('INDENT', "Expected an increase in indentation level")

    def expect_dedent(self):
        self.expect('DEDENT', "Expected a decrease in indentation level")

    def expect_newline(self, message="Expected a newline", ignore_semicolon=False):
        # Expect either a newline or end of file
        useless_trailing_semicolon = None
        if ignore_semicolon and self.sy == ';':
            useless_trailing_semicolon = self.position()
            self.next()
        if self.sy != 'EOF':
            self.expect('NEWLINE', message)
        if useless_trailing_semicolon is not None:
            warning(useless_trailing_semicolon, "useless trailing semicolon")
<commit_msg>remove unused import<commit_after># cython: infer_types=True, language_level=3, py2_import=True
#
#   Cython Scanner
#

from __future__ import absolute_import

import cython
cython.declare(make_lexicon=object, lexicon=object,
               any_string_prefix=unicode, IDENT=unicode,
               print_function=object, error=object, warning=object,
               os=object, platform=object)

import os
import platform

from .. import Utils
from ..Plex.Scanners import Scanner
from ..Plex.Errors import UnrecognizedInput
from .Errors import error, warning
from .Lexicon import any_string_prefix, make_lexicon, IDENT
from .Future import print_function

debug_scanner = 0
trace_scanner = 0
scanner_debug_flags = 0
scanner_dump_file = None

lexicon = None

def get_lexicon():
    global lexicon
    if not lexicon:
        lexicon = make_lexicon()
    return lexicon

#------------------------------------------------------------------

py_reserved_words = [
    "global", "nonlocal", "def", "class", "print", "del", "pass", "break",
    "continue", "return", "raise", "import", "exec", "try",
    "except", "finally", "while", "if", "elif", "else", "for",
    "in", "assert", "and", "or", "not", "is", "in", "lambda",
    "from", "yield", "with", "nonlocal",
]

pyx_reserved_words = py_reserved_words + [
    "include", "ctypedef", "cdef", "cpdef",
    "cimport", "DEF", "IF", "ELIF", "ELSE"
]

class Method(object):

    def __init__(self, name):
        self.name = name
        self.__name__ = name  # for Plex tracing

    def __call__(self, stream, text):
        return getattr(stream, self.name)(text)

#------------------------------------------------------------------

class CompileTimeScope(object):

    def __init__(self, outer = None):
        self.entries = {}
        self.outer = outer

    def declare(self, name, value):
        self.entries[name] = value

    def update(self, other):
        self.entries.update(other)

    def lookup_here(self, name):
        return self.entries[name]

    def __contains__(self, name):
        return name in self.entries

    def lookup(self, name):
        try:
            return self.lookup_here(name)
        except KeyError:
            outer = self.outer
            if outer:
                return outer.lookup(name)
            else:
                raise

def initial_compile_time_env():
    benv = CompileTimeScope()
    names = ('UNAME_SYSNAME', 'UNAME_NODENAME', 'UNAME_RELEASE',
        'UNAME_VERSION', 'UNAME_MACHINE')
    for name, value in zip(names, platform.uname()):
        benv.declare(name, value)
    try:
        import __builtin__ as builtins
    except ImportError:
        import builtins

    names = ('False', 'True',
             'abs', 'all', 'any', 'ascii', 'bin', 'bool', 'bytearray', 'bytes',
             'chr', 'cmp', 'complex', 'dict', 'divmod', 'enumerate', 'filter',
             'float', 'format', 'frozenset', 'hash', 'hex', 'int', 'len',
             'list', 'long', 'map', 'max', 'min', 'oct', 'ord', 'pow', 'range',
             'repr', 'reversed', 'round', 'set', 'slice', 'sorted', 'str',
             'sum', 'tuple', 'xrange', 'zip')

    for name in names:
        try:
            benv.declare(name, getattr(builtins, name))
        except AttributeError:
            # ignore, likely Py3
            pass
    denv = CompileTimeScope(benv)
    return denv

#------------------------------------------------------------------

class SourceDescriptor(object):
    """
    A SourceDescriptor should be considered immutable.
    """
    _file_type = 'pyx'

    _escaped_description = None
    _cmp_name = ''
    def __str__(self):
        assert False # To catch all places where a descriptor is used directly as a filename

    def set_file_type_from_name(self, filename):
        name, ext = os.path.splitext(filename)
        self._file_type = ext in ('.pyx', '.pxd', '.py') and ext[1:] or 'pyx'

    def is_cython_file(self):
        return self._file_type in ('pyx', 'pxd')

    def is_python_file(self):
        return self._file_type == 'py'

    def get_escaped_description(self):
        if self._escaped_description is None:
            self._escaped_description = \
                self.get_description().encode('ASCII', 'replace').decode("ASCII")
        return self._escaped_description

    def __gt__(self, other):
        # this is only used to provide some sort of order
        try:
            return self._cmp_name > other._cmp_name
        except AttributeError:
            return False

    def __lt__(self, other):
        # this is only used to provide some sort of order
        try:
            return self._cmp_name < other._cmp_name
        except AttributeError:
            return False

    def __le__(self, other):
        # this is only used to provide some sort of order
        try:
            return self._cmp_name <= other._cmp_name
        except AttributeError:
            return False

class FileSourceDescriptor(SourceDescriptor):
    """
    Represents a code source. A code source is a more generic abstraction
    for a "filename" (as sometimes the code doesn't come from a file).
    Instances of code sources are passed to Scanner.__init__ as the
    optional name argument and will be passed back when asking for
    the position()-tuple.
    """
    def __init__(self, filename, path_description=None):
        filename = Utils.decode_filename(filename)
        self.path_description = path_description or filename
        self.filename = filename
        self.set_file_type_from_name(filename)
        self._cmp_name = filename
        self._lines = {}

    def get_lines(self, encoding=None, error_handling=None):
        # we cache the lines only the second time this is called, in
        # order to save memory when they are only used once
        key = (encoding, error_handling)
        try:
            lines = self._lines[key]
            if lines is not None:
                return lines
        except KeyError:
            pass
        f = Utils.open_source_file(
            self.filename, encoding=encoding,
            error_handling=error_handling,
            # newline normalisation is costly before Py2.6
            require_normalised_newlines=False)
        try:
            lines = list(f)
        finally:
            f.close()
        if key in self._lines:
            self._lines[key] = lines
        else:
            # do not cache the first access, but remember that we
            # already read it once
            self._lines[key] = None
        return lines

    def get_description(self):
        try:
            return os.path.relpath(self.path_description)
        except ValueError:
            # path not under current directory => use complete file path
            return self.path_description

    def get_error_description(self):
        path = self.filename
        cwd = Utils.decode_filename(os.getcwd() + os.path.sep)
        if path.startswith(cwd):
            return path[len(cwd):]
        return path

    def get_filenametable_entry(self):
        return self.filename

    def __eq__(self, other):
        return isinstance(other, FileSourceDescriptor) and self.filename == other.filename

    def __hash__(self):
        return hash(self.filename)

    def __repr__(self):
        return "<FileSourceDescriptor:%s>" % self.filename

class StringSourceDescriptor(SourceDescriptor):
    """
    Instances of this class can be used instead of a filenames if the
    code originates from a string object.
    """
    filename = None

    def __init__(self, name, code):
        self.name = name
        #self.set_file_type_from_name(name)
        self.codelines = [x + "\n" for x in code.split("\n")]
        self._cmp_name = name

    def get_lines(self, encoding=None, error_handling=None):
        if not encoding:
            return self.codelines
        else:
            return [ line.encode(encoding, error_handling).decode(encoding)
                     for line in self.codelines ]

    def get_description(self):
        return self.name

    get_error_description = get_description

    def get_filenametable_entry(self):
        return "stringsource"

    def __hash__(self):
        return id(self)
        # Do not hash on the name, an identical string source should be the
        # same object (name is often defaulted in other places)
        # return hash(self.name)

    def __eq__(self, other):
        return isinstance(other, StringSourceDescriptor) and self.name == other.name

    def __repr__(self):
        return "<StringSourceDescriptor:%s>" % self.name

#------------------------------------------------------------------

class PyrexScanner(Scanner):
    #  context            Context  Compilation context
    #  included_files     [string] Files included with 'include' statement
    #  compile_time_env   dict     Environment for conditional compilation
    #  compile_time_eval  boolean  In a true conditional compilation context
    #  compile_time_expr  boolean  In a compile-time expression context

    def __init__(self, file, filename, parent_scanner=None,
                 scope=None, context=None, source_encoding=None, parse_comments=True, initial_pos=None):
        Scanner.__init__(self, get_lexicon(), file, filename, initial_pos)
        if parent_scanner:
            self.context = parent_scanner.context
            self.included_files = parent_scanner.included_files
            self.compile_time_env = parent_scanner.compile_time_env
            self.compile_time_eval = parent_scanner.compile_time_eval
            self.compile_time_expr = parent_scanner.compile_time_expr
        else:
            self.context = context
            self.included_files = scope.included_files
            self.compile_time_env = initial_compile_time_env()
            self.compile_time_eval = 1
            self.compile_time_expr = 0
            if getattr(context.options, 'compile_time_env', None):
                self.compile_time_env.update(context.options.compile_time_env)
        self.parse_comments = parse_comments
        self.source_encoding = source_encoding
        if filename.is_python_file():
            self.in_python_file = True
            self.keywords = set(py_reserved_words)
        else:
            self.in_python_file = False
            self.keywords = set(pyx_reserved_words)
        self.trace = trace_scanner
        self.indentation_stack = [0]
        self.indentation_char = None
        self.bracket_nesting_level = 0
        self.begin('INDENT')
        self.sy = ''
        self.next()

    def commentline(self, text):
        if self.parse_comments:
            self.produce('commentline', text)

    def current_level(self):
        return self.indentation_stack[-1]

    def open_bracket_action(self, text):
        self.bracket_nesting_level += 1
        return text

    def close_bracket_action(self, text):
        self.bracket_nesting_level -= 1
        return text

    def newline_action(self, text):
        if self.bracket_nesting_level == 0:
            self.begin('INDENT')
            self.produce('NEWLINE', '')

    string_states = {
        "'":   'SQ_STRING',
        '"':   'DQ_STRING',
        "'''": 'TSQ_STRING',
        '"""': 'TDQ_STRING'
    }

    def begin_string_action(self, text):
        while text[:1] in any_string_prefix:
            text = text[1:]
        self.begin(self.string_states[text])
        self.produce('BEGIN_STRING')

    def end_string_action(self, text):
        self.begin('')
        self.produce('END_STRING')

    def unclosed_string_action(self, text):
        self.end_string_action(text)
        self.error("Unclosed string literal")

    def indentation_action(self, text):
        self.begin('')
        # Indentation within brackets should be ignored.
        #if self.bracket_nesting_level > 0:
        #    return
        # Check that tabs and spaces are being used consistently.
        if text:
            c = text[0]
            #print "Scanner.indentation_action: indent with", repr(c) ###
            if self.indentation_char is None:
                self.indentation_char = c
                #print "Scanner.indentation_action: setting indent_char to", repr(c)
            else:
                if self.indentation_char != c:
                    self.error("Mixed use of tabs and spaces")
            if text.replace(c, "") != "":
                self.error("Mixed use of tabs and spaces")
        # Figure out how many indents/dedents to do
        current_level = self.current_level()
        new_level = len(text)
        #print "Changing indent level from", current_level, "to", new_level ###
        if new_level == current_level:
            return
        elif new_level > current_level:
            #print "...pushing level", new_level ###
            self.indentation_stack.append(new_level)
            self.produce('INDENT', '')
        else:
            while new_level < self.current_level():
                #print "...popping level", self.indentation_stack[-1] ###
                self.indentation_stack.pop()
                self.produce('DEDENT', '')
            #print "...current level now", self.current_level() ###
            if new_level != self.current_level():
                self.error("Inconsistent indentation")

    def eof_action(self, text):
        while len(self.indentation_stack) > 1:
            self.produce('DEDENT', '')
            self.indentation_stack.pop()
        self.produce('EOF', '')

    def next(self):
        try:
            sy, systring = self.read()
        except UnrecognizedInput:
            self.error("Unrecognized character")
            return  # just a marker, error() always raises
        if sy == IDENT:
            if systring in self.keywords:
                if systring == u'print' and print_function in self.context.future_directives:
                    self.keywords.discard('print')
                elif systring == u'exec' and self.context.language_level >= 3:
                    self.keywords.discard('exec')
                else:
                    sy = systring
            systring = self.context.intern_ustring(systring)
        self.sy = sy
        self.systring = systring
        if False: # debug_scanner:
            _, line, col = self.position()
            if not self.systring or self.sy == self.systring:
                t = self.sy
            else:
                t = "%s %s" % (self.sy, self.systring)
            print("--- %3d %2d %s" % (line, col, t))

    def peek(self):
        saved = self.sy, self.systring
        self.next()
        next = self.sy, self.systring
        self.unread(*next)
        self.sy, self.systring = saved
        return next

    def put_back(self, sy, systring):
        self.unread(self.sy, self.systring)
        self.sy = sy
        self.systring = systring

    def unread(self, token, value):
        # This method should be added to Plex
        self.queue.insert(0, (token, value))

    def error(self, message, pos=None, fatal=True):
        if pos is None:
            pos = self.position()
        if self.sy == 'INDENT':
            error(pos, "Possible inconsistent indentation")
        err = error(pos, message)
        if fatal: raise err

    def expect(self, what, message=None):
        if self.sy == what:
            self.next()
        else:
            self.expected(what, message)

    def expect_keyword(self, what, message=None):
        if self.sy == IDENT and self.systring == what:
            self.next()
        else:
            self.expected(what, message)

    def expected(self, what, message = None):
        if message:
            self.error(message)
        else:
            if self.sy == IDENT:
                found = self.systring
            else:
                found = self.sy
            self.error("Expected '%s', found '%s'" % (what, found))

    def expect_indent(self):
        self.expect('INDENT', "Expected an increase in indentation level")

    def expect_dedent(self):
        self.expect('DEDENT', "Expected a decrease in indentation level")

    def expect_newline(self, message="Expected a newline", ignore_semicolon=False):
        # Expect either a newline or end of file
        useless_trailing_semicolon = None
        if ignore_semicolon and self.sy == ';':
            useless_trailing_semicolon = self.position()
            self.next()
        if self.sy != 'EOF':
            self.expect('NEWLINE', message)
        if useless_trailing_semicolon is not None:
            warning(useless_trailing_semicolon, "useless trailing semicolon")
<|endoftext|>


<commit_before>import socket
from typing import Any, BinaryIO, Dict, Iterable, List, Optional, Tuple, Union

__version__ = '0.5.0'

Body = Union[bytes, str]
Stats = Dict[str, Union[str, int]]

DEFAULT_TUBE = 'default'
DEFAULT_PRIORITY = 2**16
DEFAULT_DELAY = 0
DEFAULT_TTR = 60

class Job:
    """A job returned from the server."""
    __slots__ = ('id', 'body')

    def __init__(self, id: int, body: Body) -> None:
        self.id = id
        self.body = body

JobOrID = Union[Job, int]

class Error(Exception):
    """Base class for non-connection related exceptions. Connection related
    issues use the built-in ``ConnectionError``.
    """

class UnknownResponseError(Error):
    """The server sent a response that this client does not understand."""

    def __init__(self, status: bytes, values: List[bytes]) -> None:
        self.status = status
        self.values = values

class BeanstalkdError(Error):
    """Base class for error messages returned from the server."""

class BadFormatError(BeanstalkdError):
    """The client sent a malformed command."""

class BuriedError(BeanstalkdError):
    """The server ran out of memory trying to grow the priority queue and had to
    bury the job.
    """

    def __init__(self, values: List[bytes] = None) -> None:
        if values:
            self.id = int(values[0])  # type: Optional[int]
        else:
            self.id = None

class DeadlineSoonError(BeanstalkdError):
    """The client has a reserved job timing out within the next second."""

class DrainingError(BeanstalkdError):
    """The client tried to insert a job while the server was in drain mode."""

class ExpectedCrlfError(BeanstalkdError):
    """The client sent a job body without a trailing CRLF."""

class InternalError(BeanstalkdError):
    """The server detected an internal error."""

class JobTooBigError(BeanstalkdError):
    """The client attempted to insert a job larger than ``max-job-size``."""

class NotFoundError(BeanstalkdError):
    """For the delete, release, bury, and kick commands, it means that the job
    does not exist or is not reserved by the client.

    For the peek commands, it means the requested job does not exist or that
    there are no jobs in the requested state.
    """

class NotIgnoredError(BeanstalkdError):
    """The client attempted to ignore the only tube on its watch list."""

class OutOfMemoryError(BeanstalkdError):
    """The server could not allocate enough memory for a job."""

class TimedOutError(BeanstalkdError):
    """A job could not be reserved within the specified timeout."""

class UnknownCommandError(BeanstalkdError):
    """The client sent a command that the server does not understand."""

ERROR_RESPONSES = {
    b'BAD_FORMAT':      BadFormatError,
    b'BURIED':          BuriedError,
    b'DEADLINE_SOON':   DeadlineSoonError,
    b'DRAINING':        DrainingError,
    b'EXPECTED_CRLF':   ExpectedCrlfError,
    b'INTERNAL_ERROR':  InternalError,
    b'JOB_TOO_BIG':     JobTooBigError,
    b'NOT_FOUND':       NotFoundError,
    b'NOT_IGNORED':     NotIgnoredError,
    b'OUT_OF_MEMORY':   OutOfMemoryError,
    b'TIMED_OUT':       TimedOutError,
    b'UNKNOWN_COMMAND': UnknownCommandError,
}

class Client:
    """A client implementing the beanstalk protocol. Upon creation a TCP
    connection with beanstalkd is established and tubes are initialized.

    :param host: The IP or hostname of the server.
    :param port: The port the server is running on.
    :param encoding: The encoding used to encode and decode job bodies.
    :param use: The tube to use after connecting.
    :param watch: The tubes to watch after connecting. The ``default`` tube will
                  be ignored if it's not included.
    """

    def __init__(self,
                 host: str = '127.0.0.1',
                 port: int = 11300,
                 encoding: Optional[str] = 'utf-8',
                 use: str = DEFAULT_TUBE,
                 watch: Union[str, Iterable[str]] = DEFAULT_TUBE) -> None:
        self._sock = socket.create_connection((host, port))
        self._reader = self._sock.makefile('rb')  # type: BinaryIO
        self.encoding = encoding

        if use != DEFAULT_TUBE:
            self.use(use)

        if isinstance(watch, str):
            if watch != DEFAULT_TUBE:
                self.watch(watch)
                self.ignore(DEFAULT_TUBE)
        else:
            for tube in watch:
                self.watch(tube)
            if DEFAULT_TUBE not in watch:
                self.ignore(DEFAULT_TUBE)

    def __enter__(self) -> 'Client':
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()

    def close(self) -> None:
        """Closes the TCP connection to beanstalkd. The client instance should
        not be used after calling this method."""
        self._reader.close()
        self._sock.close()

    def _send_cmd(self, cmd: bytes, expected: bytes) -> List[bytes]:
        self._sock.sendall(cmd + b'\r\n')
        line = self._reader.readline()
        return _parse_response(line, expected)

    def _read_chunk(self, size: int) -> bytes:
        data = self._reader.read(size + 2)
        return _parse_chunk(data, size)

    def _int_cmd(self, cmd: bytes, expected: bytes) -> int:
        n, = self._send_cmd(cmd, expected)
        return int(n)

    def _job_cmd(self, cmd: bytes, expected: bytes) -> Job:
        id, size = (int(n) for n in self._send_cmd(cmd, expected))
        chunk = self._read_chunk(size)
        if self.encoding is None:
            body = chunk  # type: Body
        else:
            body = chunk.decode(self.encoding)
        return Job(id, body)

    def _peek_cmd(self, cmd: bytes) -> Job:
        return self._job_cmd(cmd, b'FOUND')

    def _stats_cmd(self, cmd: bytes) -> Stats:
        size = self._int_cmd(cmd, b'OK')
        chunk = self._read_chunk(size)
        return _parse_simple_yaml(chunk)

    def put(self,
            body: Body,
            priority: int = DEFAULT_PRIORITY,
            delay: int = DEFAULT_DELAY,
            ttr: int = DEFAULT_TTR) -> int:
        """Inserts a job into the currently used tube and returns the job ID.

        :param body: The data representing the job.
        :param priority: An integer between 0 and 4,294,967,295 where 0 is the
                         most urgent.
        :param delay: The number of seconds to delay the job for.
        :param ttr: The maximum number of seconds a job can be reserved for
                    before timing out.
        """
        if isinstance(body, str):
            if self.encoding is None:
                raise TypeError("Unable to encode string with no encoding set")
            body = body.encode(self.encoding)
        cmd = b'put %d %d %d %d\r\n%b' % (priority, delay, ttr, len(body), body)
        return self._int_cmd(cmd, b'INSERTED')

    def use(self, tube: str) -> None:
        """Changes the currently used tube.

        :param tube: The tube to use.
        """
        self._send_cmd(b'use %b' % tube.encode('ascii'), b'USING')

    def reserve(self, timeout: Optional[int] = None) -> Job:
        """Reserves a job from a tube on the watch list, giving this client
        exclusive access to it for the TTR. Returns the reserved job.

        This blocks until a job is reserved unless a ``timeout`` is given,
        which will raise a :class:`TimedOutError <greenstalk.TimedOutError>` if
        a job cannot be reserved within that time.

        :param timeout: The maximum number of seconds to wait.
        """
        if timeout is None:
            cmd = b'reserve'
        else:
            cmd = b'reserve-with-timeout %d' % timeout
        return self._job_cmd(cmd, b'RESERVED')

    def delete(self, job: JobOrID) -> None:
        """Deletes a job.

        :param job: The job or job ID to delete.
        """
        self._send_cmd(b'delete %d' % _to_id(job), b'DELETED')

    def release(self,
                job: Job,
                priority: int = DEFAULT_PRIORITY,
                delay: int = DEFAULT_DELAY) -> None:
        """Releases a reserved job.

        :param job: The job to release.
        :param priority: An integer between 0 and 4,294,967,295 where 0 is the
                         most urgent.
        :param delay: The number of seconds to delay the job for.
        """
        self._send_cmd(b'release %d %d %d' % (job.id, priority, delay), b'RELEASED')

    def bury(self, job: Job, priority: int = DEFAULT_PRIORITY) -> None:
        """Buries a reserved job.

        :param job: The job to bury.
        :param priority: An integer between 0 and 4,294,967,295 where 0 is the
                         most urgent.
        """
        self._send_cmd(b'bury %d %d' % (job.id, priority), b'BURIED')

    def touch(self, job: Job) -> None:
        """Refreshes the TTR of a reserved job.

        :param job: The job to touch.
        """
        self._send_cmd(b'touch %d' % job.id, b'TOUCHED')

    def watch(self, tube: str) -> int:
        """Adds a tube to the watch list. Returns the number of tubes this
        client is watching.

        :param tube: The tube to watch.
        """
        return self._int_cmd(b'watch %b' % tube.encode('ascii'), b'WATCHING')

    def ignore(self, tube: str) -> int:
        """Removes a tube from the watch list. Returns the number of tubes this
        client is watching.

        :param tube: The tube to ignore.
        """
        return self._int_cmd(b'ignore %b' % tube.encode('ascii'), b'WATCHING')

    def peek(self, id: int) -> Job:
        """Returns a job by ID.

        :param id: The ID of the job to peek.
        """
        return self._peek_cmd(b'peek %d' % id)

    def peek_ready(self) -> Job:
        """Returns the next ready job in the currently used tube."""
        return self._peek_cmd(b'peek-ready')

    def peek_delayed(self) -> Job:
        """Returns the next available delayed job in the currently used tube."""
        return self._peek_cmd(b'peek-delayed')

    def peek_buried(self) -> Job:
        """Returns the oldest buried job in the currently used tube."""
        return self._peek_cmd(b'peek-buried')

    def kick(self, bound: int) -> int:
        """Moves delayed and buried jobs into the ready queue and returns the
        number of jobs effected.

        Only jobs from the currently used tube are moved.

        A kick will only move jobs in a single state. If there are any buried
        jobs, only those will be moved. Otherwise delayed jobs will be moved.

        :param bound: The maximum number of jobs to kick.
        """
        return self._int_cmd(b'kick %d' % bound, b'KICKED')

    def kick_job(self, job: JobOrID) -> None:
        """Moves a delayed or buried job into the ready queue.

        :param job: The job or job ID to kick.
        """
        self._send_cmd(b'kick-job %d' % _to_id(job), b'KICKED')

    def stats_job(self, job: JobOrID) -> Stats:
        """Returns job statistics.

        :param job: The job or job ID to return statistics for.
        """
        return self._stats_cmd(b'stats-job %d' % _to_id(job))

    def stats_tube(self, tube: str) -> Stats:
        """Returns tube statistics.

        :param tube: The tube to return statistics for.
        """
        return self._stats_cmd(b'stats-tube %b' % tube.encode('ascii'))

    def stats(self) -> Stats:
        """Returns system statistics."""
        return self._stats_cmd(b'stats')

    def pause_tube(self, tube: str, delay: int) -> None:
        """Prevents jobs from being reserved from a tube for a period of time.

        :param tube: The tube to pause.
        :param delay: The number of seconds to pause the tube for.
        """
        self._send_cmd(b'pause-tube %b %d' % (tube.encode('ascii'), delay), b'PAUSED')

def _to_id(j: JobOrID) -> int:
    return j.id if isinstance(j, Job) else j

def _parse_response(line: bytes, expected: bytes) -> List[bytes]:
    if not line:
        raise ConnectionError("Unexpected EOF")

    assert line[-2:] == b'\r\n'
    line = line[:-2]

    status, *values = line.split()

    if status == expected:
        return values

    if status in ERROR_RESPONSES:
        raise ERROR_RESPONSES[status](values)

    raise UnknownResponseError(status, values)

def _parse_chunk(data: bytes, size: int) -> bytes:
    assert data[-2:] == b'\r\n'
    data = data[:-2]

    if len(data) != size:
        raise ConnectionError("Unexpected EOF reading chunk")

    return data

def _parse_simple_yaml(buf: bytes) -> Stats:
    data = buf.decode('ascii')

    assert data[:4] == '---\n'
    data = data[4:]  # strip YAML head

    stats = {}
    for line in data.splitlines():
        key, value = line.split(': ', 1)  # type: Tuple[str, Union[str, int]]
        try:
            value = int(value)
        except ValueError:
            pass
        stats[key] = value

    return stats
<commit_msg>Make mypy happy<commit_after>import socket
from typing import Any, BinaryIO, Dict, Iterable, List, Optional, Tuple, Union

__version__ = '0.5.0'

Body = Union[bytes, str]
Stats = Dict[str, Union[str, int]]

DEFAULT_TUBE = 'default'
DEFAULT_PRIORITY = 2**16
DEFAULT_DELAY = 0
DEFAULT_TTR = 60

class Job:
    """A job returned from the server."""
    __slots__ = ('id', 'body')

    def __init__(self, id: int, body: Body) -> None:
        self.id = id
        self.body = body

JobOrID = Union[Job, int]

class Error(Exception):
    """Base class for non-connection related exceptions. Connection related
    issues use the built-in ``ConnectionError``.
    """

class UnknownResponseError(Error):
    """The server sent a response that this client does not understand."""

    def __init__(self, status: bytes, values: List[bytes]) -> None:
        self.status = status
        self.values = values

class BeanstalkdError(Error):
    """Base class for error messages returned from the server."""

class BadFormatError(BeanstalkdError):
    """The client sent a malformed command."""

class BuriedError(BeanstalkdError):
    """The server ran out of memory trying to grow the priority queue and had to
    bury the job.
    """

    def __init__(self, values: List[bytes] = None) -> None:
        if values:
            self.id = int(values[0])  # type: Optional[int]
        else:
            self.id = None

class DeadlineSoonError(BeanstalkdError):
    """The client has a reserved job timing out within the next second."""

class DrainingError(BeanstalkdError):
    """The client tried to insert a job while the server was in drain mode."""

class ExpectedCrlfError(BeanstalkdError):
    """The client sent a job body without a trailing CRLF."""

class InternalError(BeanstalkdError):
    """The server detected an internal error."""

class JobTooBigError(BeanstalkdError):
    """The client attempted to insert a job larger than ``max-job-size``."""

class NotFoundError(BeanstalkdError):
    """For the delete, release, bury, and kick commands, it means that the job
    does not exist or is not reserved by the client.

    For the peek commands, it means the requested job does not exist or that
    there are no jobs in the requested state.
    """

class NotIgnoredError(BeanstalkdError):
    """The client attempted to ignore the only tube on its watch list."""

class OutOfMemoryError(BeanstalkdError):
    """The server could not allocate enough memory for a job."""

class TimedOutError(BeanstalkdError):
    """A job could not be reserved within the specified timeout."""

class UnknownCommandError(BeanstalkdError):
    """The client sent a command that the server does not understand."""

ERROR_RESPONSES = {
    b'BAD_FORMAT':      BadFormatError,
    b'BURIED':          BuriedError,
    b'DEADLINE_SOON':   DeadlineSoonError,
    b'DRAINING':        DrainingError,
    b'EXPECTED_CRLF':   ExpectedCrlfError,
    b'INTERNAL_ERROR':  InternalError,
    b'JOB_TOO_BIG':     JobTooBigError,
    b'NOT_FOUND':       NotFoundError,
    b'NOT_IGNORED':     NotIgnoredError,
    b'OUT_OF_MEMORY':   OutOfMemoryError,
    b'TIMED_OUT':       TimedOutError,
    b'UNKNOWN_COMMAND': UnknownCommandError,
}

class Client:
    """A client implementing the beanstalk protocol. Upon creation a TCP
    connection with beanstalkd is established and tubes are initialized.

    :param host: The IP or hostname of the server.
    :param port: The port the server is running on.
    :param encoding: The encoding used to encode and decode job bodies.
    :param use: The tube to use after connecting.
    :param watch: The tubes to watch after connecting. The ``default`` tube will
                  be ignored if it's not included.
    """

    def __init__(self,
                 host: str = '127.0.0.1',
                 port: int = 11300,
                 encoding: Optional[str] = 'utf-8',
                 use: str = DEFAULT_TUBE,
                 watch: Union[str, Iterable[str]] = DEFAULT_TUBE) -> None:
        self._sock = socket.create_connection((host, port))
        self._reader = self._sock.makefile('rb')  # type: BinaryIO
        self.encoding = encoding

        if use != DEFAULT_TUBE:
            self.use(use)

        if isinstance(watch, str):
            if watch != DEFAULT_TUBE:
                self.watch(watch)
                self.ignore(DEFAULT_TUBE)
        else:
            for tube in watch:
                self.watch(tube)
            if DEFAULT_TUBE not in watch:
                self.ignore(DEFAULT_TUBE)

    def __enter__(self) -> 'Client':
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()

    def close(self) -> None:
        """Closes the TCP connection to beanstalkd. The client instance should
        not be used after calling this method."""
        self._reader.close()
        self._sock.close()

    def _send_cmd(self, cmd: bytes, expected: bytes) -> List[bytes]:
        self._sock.sendall(cmd + b'\r\n')
        line = self._reader.readline()
        return _parse_response(line, expected)

    def _read_chunk(self, size: int) -> bytes:
        data = self._reader.read(size + 2)
        return _parse_chunk(data, size)

    def _int_cmd(self, cmd: bytes, expected: bytes) -> int:
        n, = self._send_cmd(cmd, expected)
        return int(n)

    def _job_cmd(self, cmd: bytes, expected: bytes) -> Job:
        id, size = (int(n) for n in self._send_cmd(cmd, expected))
        chunk = self._read_chunk(size)
        if self.encoding is None:
            body = chunk  # type: Body
        else:
            body = chunk.decode(self.encoding)
        return Job(id, body)

    def _peek_cmd(self, cmd: bytes) -> Job:
        return self._job_cmd(cmd, b'FOUND')

    def _stats_cmd(self, cmd: bytes) -> Stats:
        size = self._int_cmd(cmd, b'OK')
        chunk = self._read_chunk(size)
        return _parse_simple_yaml(chunk)

    def put(self,
            body: Body,
            priority: int = DEFAULT_PRIORITY,
            delay: int = DEFAULT_DELAY,
            ttr: int = DEFAULT_TTR) -> int:
        """Inserts a job into the currently used tube and returns the job ID.

        :param body: The data representing the job.
        :param priority: An integer between 0 and 4,294,967,295 where 0 is the
                         most urgent.
        :param delay: The number of seconds to delay the job for.
        :param ttr: The maximum number of seconds a job can be reserved for
                    before timing out.
        """
        if isinstance(body, str):
            if self.encoding is None:
                raise TypeError("Unable to encode string with no encoding set")
            body = body.encode(self.encoding)
        cmd = b'put %d %d %d %d\r\n%b' % (priority, delay, ttr, len(body), body)
        return self._int_cmd(cmd, b'INSERTED')

    def use(self, tube: str) -> None:
        """Changes the currently used tube.

        :param tube: The tube to use.
        """
        self._send_cmd(b'use %b' % tube.encode('ascii'), b'USING')

    def reserve(self, timeout: Optional[int] = None) -> Job:
        """Reserves a job from a tube on the watch list, giving this client
        exclusive access to it for the TTR. Returns the reserved job.

        This blocks until a job is reserved unless a ``timeout`` is given,
        which will raise a :class:`TimedOutError <greenstalk.TimedOutError>` if
        a job cannot be reserved within that time.

        :param timeout: The maximum number of seconds to wait.
        """
        if timeout is None:
            cmd = b'reserve'
        else:
            cmd = b'reserve-with-timeout %d' % timeout
        return self._job_cmd(cmd, b'RESERVED')

    def delete(self, job: JobOrID) -> None:
        """Deletes a job.

        :param job: The job or job ID to delete.
        """
        self._send_cmd(b'delete %d' % _to_id(job), b'DELETED')

    def release(self,
                job: Job,
                priority: int = DEFAULT_PRIORITY,
                delay: int = DEFAULT_DELAY) -> None:
        """Releases a reserved job.

        :param job: The job to release.
        :param priority: An integer between 0 and 4,294,967,295 where 0 is the
                         most urgent.
        :param delay: The number of seconds to delay the job for.
        """
        self._send_cmd(b'release %d %d %d' % (job.id, priority, delay), b'RELEASED')

    def bury(self, job: Job, priority: int = DEFAULT_PRIORITY) -> None:
        """Buries a reserved job.

        :param job: The job to bury.
        :param priority: An integer between 0 and 4,294,967,295 where 0 is the
                         most urgent.
        """
        self._send_cmd(b'bury %d %d' % (job.id, priority), b'BURIED')

    def touch(self, job: Job) -> None:
        """Refreshes the TTR of a reserved job.

        :param job: The job to touch.
        """
        self._send_cmd(b'touch %d' % job.id, b'TOUCHED')

    def watch(self, tube: str) -> int:
        """Adds a tube to the watch list. Returns the number of tubes this
        client is watching.

        :param tube: The tube to watch.
        """
        return self._int_cmd(b'watch %b' % tube.encode('ascii'), b'WATCHING')

    def ignore(self, tube: str) -> int:
        """Removes a tube from the watch list. Returns the number of tubes this
        client is watching.

        :param tube: The tube to ignore.
        """
        return self._int_cmd(b'ignore %b' % tube.encode('ascii'), b'WATCHING')

    def peek(self, id: int) -> Job:
        """Returns a job by ID.

        :param id: The ID of the job to peek.
        """
        return self._peek_cmd(b'peek %d' % id)

    def peek_ready(self) -> Job:
        """Returns the next ready job in the currently used tube."""
        return self._peek_cmd(b'peek-ready')

    def peek_delayed(self) -> Job:
        """Returns the next available delayed job in the currently used tube."""
        return self._peek_cmd(b'peek-delayed')

    def peek_buried(self) -> Job:
        """Returns the oldest buried job in the currently used tube."""
        return self._peek_cmd(b'peek-buried')

    def kick(self, bound: int) -> int:
        """Moves delayed and buried jobs into the ready queue and returns the
        number of jobs effected.

        Only jobs from the currently used tube are moved.

        A kick will only move jobs in a single state. If there are any buried
        jobs, only those will be moved. Otherwise delayed jobs will be moved.

        :param bound: The maximum number of jobs to kick.
        """
        return self._int_cmd(b'kick %d' % bound, b'KICKED')

    def kick_job(self, job: JobOrID) -> None:
        """Moves a delayed or buried job into the ready queue.

        :param job: The job or job ID to kick.
        """
        self._send_cmd(b'kick-job %d' % _to_id(job), b'KICKED')

    def stats_job(self, job: JobOrID) -> Stats:
        """Returns job statistics.

        :param job: The job or job ID to return statistics for.
        """
        return self._stats_cmd(b'stats-job %d' % _to_id(job))

    def stats_tube(self, tube: str) -> Stats:
        """Returns tube statistics.

        :param tube: The tube to return statistics for.
        """
        return self._stats_cmd(b'stats-tube %b' % tube.encode('ascii'))

    def stats(self) -> Stats:
        """Returns system statistics."""
        return self._stats_cmd(b'stats')

    def pause_tube(self, tube: str, delay: int) -> None:
        """Prevents jobs from being reserved from a tube for a period of time.

        :param tube: The tube to pause.
        :param delay: The number of seconds to pause the tube for.
        """
        self._send_cmd(b'pause-tube %b %d' % (tube.encode('ascii'), delay), b'PAUSED')

def _to_id(j: JobOrID) -> int:
    return j.id if isinstance(j, Job) else j

def _parse_response(line: bytes, expected: bytes) -> List[bytes]:
    if not line:
        raise ConnectionError("Unexpected EOF")

    assert line[-2:] == b'\r\n'
    line = line[:-2]

    status, *values = line.split()

    if status == expected:
        return values

    if status in ERROR_RESPONSES:
        raise ERROR_RESPONSES[status](values)

    raise UnknownResponseError(status, values)

def _parse_chunk(data: bytes, size: int) -> bytes:
    assert data[-2:] == b'\r\n'
    data = data[:-2]

    if len(data) != size:
        raise ConnectionError("Unexpected EOF reading chunk")

    return data

def _parse_simple_yaml(buf: bytes) -> Stats:
    data = buf.decode('ascii')

    assert data[:4] == '---\n'
    data = data[4:]  # strip YAML head

    stats = {}
    for line in data.splitlines():
        key, value = line.split(': ', 1)
        try:
            v = int(value)  # type: Union[int, str]
        except ValueError:
            v = value
        stats[key] = v

    return stats
<|endoftext|>