<commit_before>"""
A pickler to save numpy arrays in separate .npy files.
"""

# Author: Gael Varoquaux <gael dot varoquaux at normalesup dot org>
# Copyright (c) 2009 Gael Varoquaux
# License: BSD Style, 3 clauses.

import pickle
import traceback
import sys
import os
import shutil
import tempfile
import zipfile

if sys.version_info[0] == 3:
    from pickle import _Unpickler as Unpickler
else:
    from pickle import Unpickler

###############################################################################
# Utility objects for persistence.

class NDArrayWrapper(object):
    """ An object to be persisted instead of numpy arrays.

        The only thing this object does, is store the filename in wich
        the array has been persisted.
    """
    def __init__(self, filename):
        self.filename = filename

###############################################################################
# Pickler classes

class NumpyPickler(pickle.Pickler):
    """ A pickler subclass that extracts ndarrays and saves them in .npy
        files outside of the pickle.
    """

    def __init__(self, filename):
        self._filename = filename
        self._filenames = [filename, ]
        self.file = open(filename, 'wb')
        # Count the number of npy files that we have created:
        self._npy_counter = 0
        pickle.Pickler.__init__(self, self.file,
                                protocol=pickle.HIGHEST_PROTOCOL)
        # delayed import of numpy, to avoid tight coupling
        import numpy as np
        self.np = np

    def save(self, obj):
        """ Subclass the save method, to save ndarray subclasses in npy
            files, rather than pickling them. Off course, this is a
            total abuse of the Pickler class.
        """
        if isinstance(obj, self.np.ndarray):
            self._npy_counter += 1
            try:
                filename = '%s_%02i.npy' % (self._filename,
                                            self._npy_counter)
                self._filenames.append(filename)
                self.np.save(filename, obj)
                obj = NDArrayWrapper(os.path.basename(filename))
            except:
                self._npy_counter -= 1
                # XXX: We should have a logging mechanism
                print 'Failed to save %s to .npy file:\n%s' % (
                        type(obj),
                        traceback.format_exc())
        pickle.Pickler.save(self, obj)

class NumpyUnpickler(Unpickler):
    """ A subclass of the Unpickler to unpickle our numpy pickles.
    """
    dispatch = Unpickler.dispatch.copy()

    def __init__(self, filename, mmap_mode=None):
        self._filename = filename
        self.mmap_mode = mmap_mode
        self._dirname = os.path.dirname(filename)
        self.file = open(filename, 'rb')
        Unpickler.__init__(self, self.file)
        import numpy as np
        self.np = np

    def load_build(self):
        """ This method is called to set the state of a knewly created
            object.

            We capture it to replace our place-holder objects,
            NDArrayWrapper, by the array we are interested in. We
            replace directly in the stack of pickler.
        """
        Unpickler.load_build(self)
        if isinstance(self.stack[-1], NDArrayWrapper):
            nd_array_wrapper = self.stack.pop()
            if self.np.__version__ >= '1.3':
                array = self.np.load(os.path.join(self._dirname,
                                                nd_array_wrapper.filename),
                                                mmap_mode=self.mmap_mode)
            else:
                # Numpy does not have mmap_mode before 1.3
                array = self.np.load(os.path.join(self._dirname,
                                                nd_array_wrapper.filename))
            self.stack.append(array)

    # Be careful to register our new method.
    dispatch[pickle.BUILD] = load_build

###############################################################################
# Utility functions

def dump(value, filename):
    """ Persist an arbitrary Python object into a filename, with numpy arrays
        saved as separate .npy files.

        See Also
        --------
        joblib.load : corresponding loader
        joblib.dump : function to save the object in a compressed dump
    """
    try:
        pickler = NumpyPickler(filename)
        pickler.dump(value)
    finally:
        if 'pickler' in locals() and hasattr(pickler, 'file'):
            pickler.file.flush()
            pickler.file.close()
    return pickler._filenames

def load(filename, mmap_mode=None):
    """ Reconstruct a Python object and the numpy arrays it contains from
        a persisted file.

        This function loads the numpy array files saved separately. If
        the mmap_mode argument is given, it is passed to np.save and
        arrays are loaded as memmaps. As a consequence, the reconstructed
        object might not match the original pickled object.

        See Also
        --------
        joblib.dump : function to save an object
        joblib.loadz: load from a compressed file dump
    """
    try:
        unpickler = NumpyUnpickler(filename, mmap_mode=mmap_mode)
        obj = unpickler.load()
    finally:
        if 'unpickler' in locals() and hasattr(unpickler, 'file'):
            unpickler.file.close()
    return obj

def dumpz(value, filename):
    """ Persist an arbitrary Python object into a compressed zip
        filename.

        See Also
        --------
        joblib.loadz : corresponding loader
        joblib.dump : function to save an object
     """
    kwargs = dict(compression=zipfile.ZIP_DEFLATED, mode='w')
    if sys.version_info >= (2, 5):
        kwargs['allowZip64'] = True
    dump_file = zipfile.ZipFile(filename, **kwargs)

    # Stage file in a temporary dir on disk, before writing to zip.
    tmp_dir = tempfile.mkdtemp(prefix='joblib-',
                                   dir=os.path.dirname(filename))
    try:
        dump(value, os.path.join(tmp_dir, 'joblib_dump.pkl'))
        for sub_file in os.listdir(tmp_dir):
            # We use a different arcname (archive name) to avoid having
            # the name of our tmp_dir in the archive
            dump_file.write(os.path.join(tmp_dir, sub_file),
                            arcname=os.path.join('dump_file', sub_file))
    finally:
        shutil.rmtree(tmp_dir)

    dump_file.close()
    return [filename]

def loadz(filename):
    """ Reconstruct a Python object and the numpy arrays it contains from
        a persisted zipped dump.

        See Also
        --------
        joblib.dump : function to save an object
        joblib.loadz: load from a compressed file dump
    """
    kwargs = dict(compression=zipfile.ZIP_DEFLATED, mode='r')
    if sys.version_info >= (2, 5):
        kwargs['allowZip64'] = True
    dump_file = zipfile.ZipFile(filename, **kwargs)

    # Stage files in a temporary dir on disk, before writing to zip.
    tmp_dir = tempfile.mkdtemp(prefix='joblib-',
                                   dir=os.path.dirname(filename))
    try:
        dump_file.extractall(path=tmp_dir)
        output = load(os.path.join(tmp_dir, 'dump_file', 'joblib_dump.pkl'))
    finally:
        shutil.rmtree(tmp_dir)

    dump_file.close()
    return output

<commit_msg>ENH: refactor unpickler for fast unziper<commit_after>"""
A pickler to save numpy arrays in separate .npy files.
"""

# Author: Gael Varoquaux <gael dot varoquaux at normalesup dot org>
# Copyright (c) 2009 Gael Varoquaux
# License: BSD Style, 3 clauses.

import pickle
import traceback
import sys
import os
import shutil
import tempfile
import zipfile

if sys.version_info[0] == 3:
    from pickle import _Unpickler as Unpickler
else:
    from pickle import Unpickler

###############################################################################
# Utility objects for persistence.

class NDArrayWrapper(object):
    """ An object to be persisted instead of numpy arrays.

        The only thing this object does, is store the filename in wich
        the array has been persisted.
    """
    def __init__(self, filename):
        self.filename = filename

###############################################################################
# Pickler classes

class NumpyPickler(pickle.Pickler):
    """ A pickler subclass that extracts ndarrays and saves them in .npy
        files outside of the pickle.
    """

    def __init__(self, filename):
        self._filename = filename
        self._filenames = [filename, ]
        self.file = open(filename, 'wb')
        # Count the number of npy files that we have created:
        self._npy_counter = 0
        pickle.Pickler.__init__(self, self.file,
                                protocol=pickle.HIGHEST_PROTOCOL)
        # delayed import of numpy, to avoid tight coupling
        import numpy as np
        self.np = np

    def save(self, obj):
        """ Subclass the save method, to save ndarray subclasses in npy
            files, rather than pickling them. Off course, this is a
            total abuse of the Pickler class.
        """
        if isinstance(obj, self.np.ndarray):
            self._npy_counter += 1
            try:
                filename = '%s_%02i.npy' % (self._filename,
                                            self._npy_counter)
                self._filenames.append(filename)
                self.np.save(filename, obj)
                obj = NDArrayWrapper(os.path.basename(filename))
            except:
                self._npy_counter -= 1
                # XXX: We should have a logging mechanism
                print 'Failed to save %s to .npy file:\n%s' % (
                        type(obj),
                        traceback.format_exc())
        pickle.Pickler.save(self, obj)

class NumpyUnpickler(Unpickler):
    """ A subclass of the Unpickler to unpickle our numpy pickles.
    """
    dispatch = Unpickler.dispatch.copy()

    def __init__(self, filename, mmap_mode=None):
        self._filename = os.path.basename(filename)
        self.mmap_mode = mmap_mode
        self._dirname = os.path.dirname(filename)
        file_handle = self._open_file(self._filename)
        if isinstance(file_handle, basestring):
            # To hanlde memmap, we need to have file names
            file_handle = open(file_handle, 'rb')
        self.file_handle = file_handle
        Unpickler.__init__(self, self.file_handle)
        import numpy as np
        self.np = np

    def _open_file(self, name):
        "Return the path of the given file in our store"
        return os.path.join(self._dirname, name)

    def load_build(self):
        """ This method is called to set the state of a knewly created
            object.

            We capture it to replace our place-holder objects,
            NDArrayWrapper, by the array we are interested in. We
            replace directly in the stack of pickler.
        """
        Unpickler.load_build(self)
        if isinstance(self.stack[-1], NDArrayWrapper):
            nd_array_wrapper = self.stack.pop()
            if self.np.__version__ >= '1.3':
                array = self.np.load(
                                self._open_file(nd_array_wrapper.filename),
                                mmap_mode=self.mmap_mode)
            else:
                # Numpy does not have mmap_mode before 1.3
                array = self.np.load(
                                self._open_file(nd_array_wrapper.filename),
                                mmap_mode=self.mmap_mode)
            self.stack.append(array)

    # Be careful to register our new method.
    dispatch[pickle.BUILD] = load_build

###############################################################################
# Utility functions

def dump(value, filename):
    """ Persist an arbitrary Python object into a filename, with numpy arrays
        saved as separate .npy files.

        See Also
        --------
        joblib.load : corresponding loader
        joblib.dump : function to save the object in a compressed dump
    """
    try:
        pickler = NumpyPickler(filename)
        pickler.dump(value)
    finally:
        if 'pickler' in locals() and hasattr(pickler, 'file'):
            pickler.file.flush()
            pickler.file.close()
    return pickler._filenames

def load(filename, mmap_mode=None):
    """ Reconstruct a Python object and the numpy arrays it contains from
        a persisted file.

        This function loads the numpy array files saved separately. If
        the mmap_mode argument is given, it is passed to np.save and
        arrays are loaded as memmaps. As a consequence, the reconstructed
        object might not match the original pickled object.

        See Also
        --------
        joblib.dump : function to save an object
        joblib.loadz: load from a compressed file dump
    """
    try:
        unpickler = NumpyUnpickler(filename, mmap_mode=mmap_mode)
        obj = unpickler.load()
    finally:
        if 'unpickler' in locals() and hasattr(unpickler, 'file'):
            unpickler.file.close()
    return obj

def dumpz(value, filename):
    """ Persist an arbitrary Python object into a compressed zip
        filename.

        See Also
        --------
        joblib.loadz : corresponding loader
        joblib.dump : function to save an object
     """
    kwargs = dict(compression=zipfile.ZIP_DEFLATED, mode='w')
    if sys.version_info >= (2, 5):
        kwargs['allowZip64'] = True
    dump_file = zipfile.ZipFile(filename, **kwargs)

    # Stage file in a temporary dir on disk, before writing to zip.
    tmp_dir = tempfile.mkdtemp(prefix='joblib-',
                                   dir=os.path.dirname(filename))
    try:
        dump(value, os.path.join(tmp_dir, 'joblib_dump.pkl'))
        for sub_file in os.listdir(tmp_dir):
            # We use a different arcname (archive name) to avoid having
            # the name of our tmp_dir in the archive
            dump_file.write(os.path.join(tmp_dir, sub_file),
                            arcname=os.path.join('dump_file', sub_file))
    finally:
        shutil.rmtree(tmp_dir)

    dump_file.close()
    return [filename]

def loadz(filename):
    """ Reconstruct a Python object and the numpy arrays it contains from
        a persisted zipped dump.

        See Also
        --------
        joblib.dump : function to save an object
        joblib.loadz: load from a compressed file dump
    """
    kwargs = dict(compression=zipfile.ZIP_DEFLATED, mode='r')
    if sys.version_info >= (2, 5):
        kwargs['allowZip64'] = True
    dump_file = zipfile.ZipFile(filename, **kwargs)

    # Stage files in a temporary dir on disk, before writing to zip.
    tmp_dir = tempfile.mkdtemp(prefix='joblib-',
                                   dir=os.path.dirname(filename))
    try:
        dump_file.extractall(path=tmp_dir)
        output = load(os.path.join(tmp_dir, 'dump_file', 'joblib_dump.pkl'))
    finally:
        shutil.rmtree(tmp_dir)

    dump_file.close()
    return output

<|endoftext|>


<commit_before>class Turma(object):
	def __init__(self,turma):
		self.id = turma.getId()
		self.letra = turma.getLetra()
<commit_msg>Update Turma.py<commit_after>class Turma(object):
	def __init__(self,turma):
		self.id = turma.getId()
		self.letra = turma.getLetra()
		self.id_disciplina = turma.getId_disciplina()
<|endoftext|>


<commit_before>#!/usr/bin/python

# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""This module uprevs a given package's ebuild to the next revision."""

import optparse
import os
import subprocess
import sys

import constants
if __name__ == '__main__':
  sys.path.insert(0, constants.SOURCE_ROOT)

from chromite.buildbot import portage_utilities
from chromite.lib import cros_build_lib

# TODO(sosa): Remove during OO refactor.
VERBOSE = False

# Dictionary of valid commands with usage information.
COMMAND_DICTIONARY = {
                        'commit':
                          'Marks given ebuilds as stable locally',
                        'push':
                          'Pushes previous marking of ebuilds to remote repo',
                      }

# ======================= Global Helper Functions ========================

def _Print(message):
  """Verbose print function."""
  if VERBOSE:
    cros_build_lib.Info(message)

def CleanStalePackages(boards, package_atoms):
  """Cleans up stale package info from a previous build.
  Args:
    boards: Boards to clean the packages from.
    package_atoms: The actual package atom to unmerge.
  """
  if package_atoms:
    cros_build_lib.Info('Cleaning up stale packages %s.' % package_atoms)
    for board in boards:
      unmerge_board_cmd = ['emerge-%s' % board, '--unmerge']
      unmerge_board_cmd.extend(package_atoms)
      cros_build_lib.RunCommand(unmerge_board_cmd)

    unmerge_host_cmd = ['emerge', '--unmerge']
    unmerge_host_cmd.extend(package_atoms)
    cros_build_lib.SudoRunCommand(unmerge_host_cmd)

  for board in boards:
    cros_build_lib.RunCommand(['eclean-%s' % board, '-d', 'packages'],
                              redirect_stderr=True)
  cros_build_lib.SudoRunCommand(['eclean', '-d', 'packages'],
                            redirect_stderr=True)

def _DoWeHaveLocalCommits(stable_branch, tracking_branch):
  """Returns true if there are local commits."""
  current_branch = _SimpleRunCommand('git branch | grep \*').split()[1]
  if current_branch == stable_branch:
    current_commit_id = _SimpleRunCommand('git rev-parse HEAD')
    tracking_commit_id = _SimpleRunCommand('git rev-parse %s' % tracking_branch)
    return current_commit_id != tracking_commit_id
  else:
    return False

def _CheckSaneArguments(package_list, command, options):
  """Checks to make sure the flags are sane.  Dies if arguments are not sane."""
  if not command in COMMAND_DICTIONARY.keys():
    _PrintUsageAndDie('%s is not a valid command' % command)
  if not options.packages and command == 'commit' and not options.all:
    _PrintUsageAndDie('Please specify at least one package')
  if not options.boards and command == 'commit':
    _PrintUsageAndDie('Please specify a board')
  if not os.path.isdir(options.srcroot):
    _PrintUsageAndDie('srcroot is not a valid path')
  options.srcroot = os.path.abspath(options.srcroot)

def _PrintUsageAndDie(error_message=''):
  """Prints optional error_message the usage and returns an error exit code."""
  command_usage = 'Commands: \n'
  # Add keys and usage information from dictionary.
  commands = sorted(COMMAND_DICTIONARY.keys())
  for command in commands:
    command_usage += '  %s: %s\n' % (command, COMMAND_DICTIONARY[command])
  commands_str = '|'.join(commands)
  cros_build_lib.Warning('Usage: %s FLAGS [%s]\n\n%s' % (
      sys.argv[0], commands_str, command_usage))
  if error_message:
    cros_build_lib.Die(error_message)
  else:
    sys.exit(1)

def _SimpleRunCommand(command):
  """Runs a shell command and returns stdout back to caller."""
  _Print('  + %s' % command)
  proc_handle = subprocess.Popen(command, stdout=subprocess.PIPE, shell=True)
  stdout = proc_handle.communicate()[0]
  retcode = proc_handle.wait()
  if retcode != 0:
    _Print(stdout)
    raise subprocess.CalledProcessError(retcode, command)
  return stdout

# ======================= End Global Helper Functions ========================

def PushChange(stable_branch, tracking_branch, dryrun, cwd='.'):
  """Pushes commits in the stable_branch to the remote git repository.

  Pushes local commits from calls to CommitChange to the remote git
  repository specified by current working directory. If changes are
  found to commit, they will be merged to the merge branch and pushed.
  In that case, the local repository will be left on the merge branch.

  Args:
    stable_branch: The local branch with commits we want to push.
    tracking_branch: The tracking branch of the local branch.
    dryrun: Use git push --dryrun to emulate a push.
  Raises:
      OSError: Error occurred while pushing.
  """
  if not _DoWeHaveLocalCommits(stable_branch, tracking_branch):
    cros_build_lib.Info('No work found to push.  Exiting')
    return

  # For the commit queue, our local branch may contain commits that were
  # just tested and pushed during the CommitQueueCompletion stage. Sync
  # and rebase our local branch on top of the remote commits.
  remote, push_branch = cros_build_lib.GetPushBranch(cwd)
  cros_build_lib.SyncPushBranch(cwd, remote, push_branch)

  # Check whether any local changes remain after the sync.
  if not _DoWeHaveLocalCommits(stable_branch, '%s/%s' % (remote, push_branch)):
    cros_build_lib.Info('All changes already pushed. Exiting')
    return

  description = _SimpleRunCommand('git log --format=format:%s%n%n%b ' +
      '%s/%s..%s' % (remote, push_branch, stable_branch))
  description = 'Marking set of ebuilds as stable\n\n%s' % description
  cros_build_lib.Info('Using description %s' % description)
  cros_build_lib.CreatePushBranch(constants.MERGE_BRANCH, cwd)
  _SimpleRunCommand('git merge --squash %s' % stable_branch)
  cros_build_lib.RunCommand(['git', 'commit', '-m', description])
  _SimpleRunCommand('git config push.default tracking')
  cros_build_lib.GitPushWithRetry(constants.MERGE_BRANCH, cwd=cwd,
                                  dryrun=dryrun)

class GitBranch(object):
  """Wrapper class for a git branch."""

  def __init__(self, branch_name, tracking_branch):
    """Sets up variables but does not create the branch."""
    self.branch_name = branch_name
    self.tracking_branch = tracking_branch

  def CreateBranch(self):
    GitBranch.Checkout(self)

  @classmethod
  def Checkout(cls, target):
    """Function used to check out to another GitBranch."""
    if target.branch_name == target.tracking_branch or target.Exists():
      git_cmd = 'git checkout %s -f' % target.branch_name
    else:
      git_cmd = 'repo start %s .' % target.branch_name
    _SimpleRunCommand(git_cmd)

  def Exists(self):
    """Returns True if the branch exists."""
    branch_cmd = 'git branch'
    branches = _SimpleRunCommand(branch_cmd)
    return self.branch_name in branches.split()

  def Delete(self):
    """Deletes the branch and returns the user to the master branch.

    Returns True on success.
    """
    tracking_branch = GitBranch(self.tracking_branch, self.tracking_branch)
    GitBranch.Checkout(tracking_branch)
    delete_cmd = 'repo abandon %s' % self.branch_name
    _SimpleRunCommand(delete_cmd)

def main():
  parser = optparse.OptionParser('cros_mark_as_stable OPTIONS packages')
  parser.add_option('--all', action='store_true',
                    help='Mark all packages as stable.')
  parser.add_option('-b', '--boards',
                    help='Colon-separated list of boards')
  parser.add_option('--drop_file',
                    help='File to list packages that were revved.')
  parser.add_option('--dryrun', action='store_true',
                    help='Passes dry-run to git push if pushing a change.')
  parser.add_option('-o', '--overlays',
                    help='Colon-separated list of overlays to modify.')
  parser.add_option('-p', '--packages',
                    help='Colon separated list of packages to rev.')
  parser.add_option('-r', '--srcroot',
                    default='%s/trunk/src' % os.environ['HOME'],
                    help='Path to root src directory.')
  parser.add_option('--verbose', action='store_true',
                    help='Prints out debug info.')
  (options, args) = parser.parse_args()

  global VERBOSE
  VERBOSE = options.verbose
  portage_utilities.EBuild.VERBOSE = options.verbose

  if len(args) != 1:
    _PrintUsageAndDie('Must specify a valid command [commit, push]')

  command = args[0]
  package_list = None
  if options.packages:
    package_list = options.packages.split(':')

  _CheckSaneArguments(package_list, command, options)
  if options.overlays:
    overlays = {}
    for path in options.overlays.split(':'):
      if not os.path.isdir(path):
        cros_build_lib.Die('Cannot find overlay: %s' % path)
      overlays[path] = []
  else:
    cros_build_lib.Warning('Missing --overlays argument')
    overlays = {
      '%s/private-overlays/chromeos-overlay' % options.srcroot: [],
      '%s/third_party/chromiumos-overlay' % options.srcroot: []
    }

  if command == 'commit':
    portage_utilities.BuildEBuildDictionary(
      overlays, options.all, package_list)

  tracking_branch = cros_build_lib.GetManifestDefaultBranch(options.srcroot)
  tracking_branch = 'remotes/m/' + tracking_branch

  # Contains the array of packages we actually revved.
  revved_packages = []
  new_package_atoms = []

  for overlay, ebuilds in overlays.items():
    if not os.path.isdir(overlay):
      cros_build_lib.Warning("Skipping %s" % overlay)
      continue

    # TODO(davidjames): Currently, all code that interacts with git depends on
    # the cwd being set to the overlay directory. We should instead pass in
    # this parameter so that we don't need to modify the cwd globally.
    os.chdir(overlay)

    if command == 'push':
      PushChange(constants.STABLE_EBUILD_BRANCH, tracking_branch,
                 options.dryrun)
    elif command == 'commit' and ebuilds:
      existing_branch = cros_build_lib.GetCurrentBranch('.')
      work_branch = GitBranch(constants.STABLE_EBUILD_BRANCH, tracking_branch)
      work_branch.CreateBranch()
      if not work_branch.Exists():
        cros_build_lib.Die('Unable to create stabilizing branch in %s' %
                           overlay)

      # In the case of uprevving overlays that have patches applied to them,
      # include the patched changes in the stabilizing branch.
      if existing_branch:
        _SimpleRunCommand('git rebase %s' % existing_branch)

      for ebuild in ebuilds:
        try:
          _Print('Working on %s' % ebuild.package)
          new_package = ebuild.RevWorkOnEBuild(options.srcroot)
          if new_package:
            revved_packages.append(ebuild.package)
            new_package_atoms.append('=%s' % new_package)
        except (OSError, IOError):
          cros_build_lib.Warning('Cannot rev %s\n' % ebuild.package +
                  'Note you will have to go into %s '
                  'and reset the git repo yourself.' % overlay)
          raise

  if command == 'commit':
    CleanStalePackages(options.boards.split(':'), new_package_atoms)
    if options.drop_file:
      fh = open(options.drop_file, 'w')
      fh.write(' '.join(revved_packages))
      fh.close()

if __name__ == '__main__':
  main()
<commit_msg>cros_mark_as_stable: parallelize CleanStalePackages<commit_after>#!/usr/bin/python

# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""This module uprevs a given package's ebuild to the next revision."""

import optparse
import os
import subprocess
import sys

import constants
if __name__ == '__main__':
  sys.path.insert(0, constants.SOURCE_ROOT)

from chromite.buildbot import cbuildbot_background as background
from chromite.buildbot import portage_utilities
from chromite.lib import cros_build_lib

# TODO(sosa): Remove during OO refactor.
VERBOSE = False

# Dictionary of valid commands with usage information.
COMMAND_DICTIONARY = {
                        'commit':
                          'Marks given ebuilds as stable locally',
                        'push':
                          'Pushes previous marking of ebuilds to remote repo',
                      }

# ======================= Global Helper Functions ========================

def _Print(message):
  """Verbose print function."""
  if VERBOSE:
    cros_build_lib.Info(message)

def CleanStalePackages(boards, package_atoms):
  """Cleans up stale package info from a previous build.
  Args:
    boards: Boards to clean the packages from.
    package_atoms: The actual package atom to unmerge.
  """
  if package_atoms:
    cros_build_lib.Info('Cleaning up stale packages %s.' % package_atoms)

  # First unmerge all the packages for a board, then eclean it.
  # We need these two steps to run in order (unmerge/eclean),
  # but we can let all the boards run in parallel.
  def _CleanStalePackages(board):
    if board:
      suffix = '-' + board
      runcmd = cros_build_lib.RunCommand
    else:
      suffix = ''
      runcmd = cros_build_lib.SudoRunCommand

    if package_atoms:
      runcmd(['emerge' + suffix, '-q', '--unmerge'] + package_atoms);
    runcmd(['eclean' + suffix, '-d', 'packages'],
           redirect_stdout=True, redirect_stderr=True)

  tasks = []
  for board in boards:
    tasks.append([board])
  tasks.append([None])

  background.RunTasksInProcessPool(_CleanStalePackages, tasks)

def _DoWeHaveLocalCommits(stable_branch, tracking_branch):
  """Returns true if there are local commits."""
  current_branch = _SimpleRunCommand('git branch | grep \*').split()[1]
  if current_branch == stable_branch:
    current_commit_id = _SimpleRunCommand('git rev-parse HEAD')
    tracking_commit_id = _SimpleRunCommand('git rev-parse %s' % tracking_branch)
    return current_commit_id != tracking_commit_id
  else:
    return False

def _CheckSaneArguments(package_list, command, options):
  """Checks to make sure the flags are sane.  Dies if arguments are not sane."""
  if not command in COMMAND_DICTIONARY.keys():
    _PrintUsageAndDie('%s is not a valid command' % command)
  if not options.packages and command == 'commit' and not options.all:
    _PrintUsageAndDie('Please specify at least one package')
  if not options.boards and command == 'commit':
    _PrintUsageAndDie('Please specify a board')
  if not os.path.isdir(options.srcroot):
    _PrintUsageAndDie('srcroot is not a valid path')
  options.srcroot = os.path.abspath(options.srcroot)

def _PrintUsageAndDie(error_message=''):
  """Prints optional error_message the usage and returns an error exit code."""
  command_usage = 'Commands: \n'
  # Add keys and usage information from dictionary.
  commands = sorted(COMMAND_DICTIONARY.keys())
  for command in commands:
    command_usage += '  %s: %s\n' % (command, COMMAND_DICTIONARY[command])
  commands_str = '|'.join(commands)
  cros_build_lib.Warning('Usage: %s FLAGS [%s]\n\n%s' % (
      sys.argv[0], commands_str, command_usage))
  if error_message:
    cros_build_lib.Die(error_message)
  else:
    sys.exit(1)

def _SimpleRunCommand(command):
  """Runs a shell command and returns stdout back to caller."""
  _Print('  + %s' % command)
  proc_handle = subprocess.Popen(command, stdout=subprocess.PIPE, shell=True)
  stdout = proc_handle.communicate()[0]
  retcode = proc_handle.wait()
  if retcode != 0:
    _Print(stdout)
    raise subprocess.CalledProcessError(retcode, command)
  return stdout

# ======================= End Global Helper Functions ========================

def PushChange(stable_branch, tracking_branch, dryrun, cwd='.'):
  """Pushes commits in the stable_branch to the remote git repository.

  Pushes local commits from calls to CommitChange to the remote git
  repository specified by current working directory. If changes are
  found to commit, they will be merged to the merge branch and pushed.
  In that case, the local repository will be left on the merge branch.

  Args:
    stable_branch: The local branch with commits we want to push.
    tracking_branch: The tracking branch of the local branch.
    dryrun: Use git push --dryrun to emulate a push.
  Raises:
      OSError: Error occurred while pushing.
  """
  if not _DoWeHaveLocalCommits(stable_branch, tracking_branch):
    cros_build_lib.Info('No work found to push.  Exiting')
    return

  # For the commit queue, our local branch may contain commits that were
  # just tested and pushed during the CommitQueueCompletion stage. Sync
  # and rebase our local branch on top of the remote commits.
  remote, push_branch = cros_build_lib.GetPushBranch(cwd)
  cros_build_lib.SyncPushBranch(cwd, remote, push_branch)

  # Check whether any local changes remain after the sync.
  if not _DoWeHaveLocalCommits(stable_branch, '%s/%s' % (remote, push_branch)):
    cros_build_lib.Info('All changes already pushed. Exiting')
    return

  description = _SimpleRunCommand('git log --format=format:%s%n%n%b ' +
      '%s/%s..%s' % (remote, push_branch, stable_branch))
  description = 'Marking set of ebuilds as stable\n\n%s' % description
  cros_build_lib.Info('Using description %s' % description)
  cros_build_lib.CreatePushBranch(constants.MERGE_BRANCH, cwd)
  _SimpleRunCommand('git merge --squash %s' % stable_branch)
  cros_build_lib.RunCommand(['git', 'commit', '-m', description])
  _SimpleRunCommand('git config push.default tracking')
  cros_build_lib.GitPushWithRetry(constants.MERGE_BRANCH, cwd=cwd,
                                  dryrun=dryrun)

class GitBranch(object):
  """Wrapper class for a git branch."""

  def __init__(self, branch_name, tracking_branch):
    """Sets up variables but does not create the branch."""
    self.branch_name = branch_name
    self.tracking_branch = tracking_branch

  def CreateBranch(self):
    GitBranch.Checkout(self)

  @classmethod
  def Checkout(cls, target):
    """Function used to check out to another GitBranch."""
    if target.branch_name == target.tracking_branch or target.Exists():
      git_cmd = 'git checkout %s -f' % target.branch_name
    else:
      git_cmd = 'repo start %s .' % target.branch_name
    _SimpleRunCommand(git_cmd)

  def Exists(self):
    """Returns True if the branch exists."""
    branch_cmd = 'git branch'
    branches = _SimpleRunCommand(branch_cmd)
    return self.branch_name in branches.split()

  def Delete(self):
    """Deletes the branch and returns the user to the master branch.

    Returns True on success.
    """
    tracking_branch = GitBranch(self.tracking_branch, self.tracking_branch)
    GitBranch.Checkout(tracking_branch)
    delete_cmd = 'repo abandon %s' % self.branch_name
    _SimpleRunCommand(delete_cmd)

def main():
  parser = optparse.OptionParser('cros_mark_as_stable OPTIONS packages')
  parser.add_option('--all', action='store_true',
                    help='Mark all packages as stable.')
  parser.add_option('-b', '--boards',
                    help='Colon-separated list of boards')
  parser.add_option('--drop_file',
                    help='File to list packages that were revved.')
  parser.add_option('--dryrun', action='store_true',
                    help='Passes dry-run to git push if pushing a change.')
  parser.add_option('-o', '--overlays',
                    help='Colon-separated list of overlays to modify.')
  parser.add_option('-p', '--packages',
                    help='Colon separated list of packages to rev.')
  parser.add_option('-r', '--srcroot',
                    default='%s/trunk/src' % os.environ['HOME'],
                    help='Path to root src directory.')
  parser.add_option('--verbose', action='store_true',
                    help='Prints out debug info.')
  (options, args) = parser.parse_args()

  global VERBOSE
  VERBOSE = options.verbose
  portage_utilities.EBuild.VERBOSE = options.verbose

  if len(args) != 1:
    _PrintUsageAndDie('Must specify a valid command [commit, push]')

  command = args[0]
  package_list = None
  if options.packages:
    package_list = options.packages.split(':')

  _CheckSaneArguments(package_list, command, options)
  if options.overlays:
    overlays = {}
    for path in options.overlays.split(':'):
      if not os.path.isdir(path):
        cros_build_lib.Die('Cannot find overlay: %s' % path)
      overlays[path] = []
  else:
    cros_build_lib.Warning('Missing --overlays argument')
    overlays = {
      '%s/private-overlays/chromeos-overlay' % options.srcroot: [],
      '%s/third_party/chromiumos-overlay' % options.srcroot: []
    }

  if command == 'commit':
    portage_utilities.BuildEBuildDictionary(
      overlays, options.all, package_list)

  tracking_branch = cros_build_lib.GetManifestDefaultBranch(options.srcroot)
  tracking_branch = 'remotes/m/' + tracking_branch

  # Contains the array of packages we actually revved.
  revved_packages = []
  new_package_atoms = []

  for overlay, ebuilds in overlays.items():
    if not os.path.isdir(overlay):
      cros_build_lib.Warning("Skipping %s" % overlay)
      continue

    # TODO(davidjames): Currently, all code that interacts with git depends on
    # the cwd being set to the overlay directory. We should instead pass in
    # this parameter so that we don't need to modify the cwd globally.
    os.chdir(overlay)

    if command == 'push':
      PushChange(constants.STABLE_EBUILD_BRANCH, tracking_branch,
                 options.dryrun)
    elif command == 'commit' and ebuilds:
      existing_branch = cros_build_lib.GetCurrentBranch('.')
      work_branch = GitBranch(constants.STABLE_EBUILD_BRANCH, tracking_branch)
      work_branch.CreateBranch()
      if not work_branch.Exists():
        cros_build_lib.Die('Unable to create stabilizing branch in %s' %
                           overlay)

      # In the case of uprevving overlays that have patches applied to them,
      # include the patched changes in the stabilizing branch.
      if existing_branch:
        _SimpleRunCommand('git rebase %s' % existing_branch)

      for ebuild in ebuilds:
        try:
          _Print('Working on %s' % ebuild.package)
          new_package = ebuild.RevWorkOnEBuild(options.srcroot)
          if new_package:
            revved_packages.append(ebuild.package)
            new_package_atoms.append('=%s' % new_package)
        except (OSError, IOError):
          cros_build_lib.Warning('Cannot rev %s\n' % ebuild.package +
                  'Note you will have to go into %s '
                  'and reset the git repo yourself.' % overlay)
          raise

  if command == 'commit':
    CleanStalePackages(options.boards.split(':'), new_package_atoms)
    if options.drop_file:
      fh = open(options.drop_file, 'w')
      fh.write(' '.join(revved_packages))
      fh.close()

if __name__ == '__main__':
  main()
<|endoftext|>


<commit_before>import urlparse
import os

from buildbot.config import error as config_error

# TBD use plugins!
from buildbot.config import BuilderConfig
from buildbot.schedulers.forcesched import StringParameter, CodebaseParameter
from buildbot.schedulers.triggerable import Triggerable
from buildbot.worker import Worker
from buildbot.worker import AbstractLatentWorker
from buildbot.process import factory
from buildbot.plugins import util
from buildbot import getVersion
from .important import ImportantManager
from .vcs import addRepository, getSupportedVCSTypes
from .steps import TravisSetupSteps
from .steps import TravisTrigger
from yaml import safe_load

import buildbot_travis

class TravisConfigurator(object):

    def __init__(self, config, vardir, latentRunners=False):
        self.config = config
        self.vardir = vardir
        self.latentRunners = latentRunners
        self.passwords = {}
        self.properties = {}
        self.repositories = {}
        self.cfgdict = {}
        self.importantManager = None
        self.change_hook_dialects = {}
        config.setdefault("builders", [])
        config.setdefault("schedulers", [])
        config.setdefault("change_source", [])
        config.setdefault("services", [])
        config.setdefault("status", [])
        self.defaultEnv = {}
        self.defaultStages = []
        # we are not really multimaster, but this remove some checks
        config.setdefault("multiMaster", True)
        config['codebaseGenerator'] = lambda chdict: chdict['project']
        self.config['title'] = os.environ.get(
            'buildbotTitle', "buildbot travis")

    def add_password(self, scheme, netloc, username, password):
        self.passwords[(scheme, netloc)] = (username, password)

    def fromYaml(self, path):
        buildbot_travis.api.setYamlPath(path)
        with open(path) as f:
            y = safe_load(f)

        return self.fromDict(y)

    def fromDict(self, y):
        buildbot_travis.api.setCfg(y)
        self.cfgdict = y
        self.importantManager = ImportantManager(
            y.get("not_important_files", []))
        self.defaultEnv = y.get("env", {})
        for k, v in self.defaultEnv.items():
            if not (isinstance(v, list) or isinstance(v, basestring)):
                config_error(
                    "'env' values must be strings or lists ; key %s is incorrect: %s" % (k, type(v)))
        for p in y.get("projects", []):
            self.define_travis_builder(**p)
        self.defaultStages = y.get("stages", [])
        for s in self.defaultStages:
            if not isinstance(s, basestring):
                config_error(
                    "'stages' values must be strings ; stage %s is incorrect: %s" % (s, type(s)))

        PORT = int(os.environ.get('PORT', 8020))
        self.config['buildbotURL'] = os.environ.get(
            'buildbotURL', "http://localhost:%d/" % (PORT, ))
        # minimalistic config to activate new web UI
        self.config['www'] = dict(port=PORT,
                                  change_hook_dialects=self.change_hook_dialects,
                                  plugins=dict(buildbot_travis={'cfg': self.cfgdict,
                                                                'supported_vcs': getSupportedVCSTypes()}),
                                  versions=[('buildbot_travis', getVersion(__file__))])

    def fromDb(self):
        buildbot_travis.api.useDbConfig()
        dbConfig = util.DbConfig(self.config, self.vardir)
        return self.fromDict(dbConfig.get("travis", {}))

    def get_spawner_workers(self):
        workers = [s.workername for s in self.config[
            'workers'] if isinstance(s, Worker)]
        return workers

    def get_runner_workers(self):
        if self.latentRunners:
            WorkerClass = AbstractLatentWorker
        else:
            WorkerClass = Worker
        workers = [s.workername for s in self.config[
            'workers'] if isinstance(s, WorkerClass)]
        return workers

    def define_travis_builder(self, name, repository, tags=None, **kwargs):
        name = str(name)
        repository = str(repository)
        job_name = "%s-job" % name
        try_name = "%s-try" % name
        deploy_name = "%s-deploy" % name
        spawner_name = name
        if tags is None:
            tags = []

        def formatTag(tag):
            if isinstance(tag, basestring):
                return str(tag)
            return str(tag['text'])

        tags = map(formatTag, tags)
        if 'username' not in kwargs and 'password' not in kwargs:
            p = urlparse.urlparse(repository)
            k = (p.scheme, p.netloc)
            if k in self.passwords:
                kwargs['username'], kwargs['password'] = self.passwords[k]

        codebases = {spawner_name: {'repository': repository}}
        for subrepo in kwargs.get('subrepos', []):
            codebases[subrepo['project']] = {
                'repository': subrepo['repository']}

        vcsManager = addRepository(
            name, dict(name=name, repository=repository, **kwargs))
        vcsManager.vardir = self.vardir

        # Define the builder for the main job
        f = factory.BuildFactory()
        vcsManager.addSourceSteps(f)
        f.addStep(TravisSetupSteps())

        self.config['builders'].append(BuilderConfig(
            name=job_name,
            workernames=self.get_runner_workers(),
            properties=self.properties,
            collapseRequests=False,
            env=self.defaultEnv,
            tags=["job", name] + tags,
            factory=f
        ))

        self.config['schedulers'].append(Triggerable(
            name=job_name,
            builderNames=[job_name],
            codebases=codebases,
        ))

        # Define the builder for a spawner
        f = factory.BuildFactory()
        vcsManager.addSourceSteps(f)
        f.addStep(TravisTrigger(
            scheduler=job_name,
        ))
        properties = dict(TRAVIS_PULL_REQUEST=False)
        properties.update(self.properties)
        self.config['builders'].append(BuilderConfig(
            name=spawner_name,
            workernames=self.get_spawner_workers(),
            properties=properties,
            tags=["trunk", name] + tags,
            factory=f
        ))

        # Define the builder for the deployment of the project
        f = factory.BuildFactory()
        vcsManager.addSourceSteps(f)
        f.addStep(TravisSetupSteps())

        # To manage deployment properly (with change traceability),
        # we need the version and the target deployment environment or "stage"
        version = StringParameter(name='version', label='GIT tag',
                                  hide=False, required=False, size=20)
        stage = StringParameter(name='stage', label='Stage',
                                hide=False, required=False,size=20)

        dep_properties = [version, stage]

        self.config['builders'].append(BuilderConfig(
            name=deploy_name,
            workernames=self.get_runner_workers(),
            tags=["deploy", name] + tags,
            factory=f
        ))

        ))

        if vcsManager.supportsTry:
            properties = dict(TRAVIS_PULL_REQUEST=True)
            properties.update(self.properties)
            # Define the builder for try job
            f = factory.BuildFactory()
            vcsManager.addSourceSteps(f)
            f.addStep(TravisTrigger(
                scheduler=job_name,
            ))

            self.config['builders'].append(BuilderConfig(
                name=try_name,
                workernames=self.get_spawner_workers(),
                properties=properties,
                tags=["try", name] + tags,
                factory=f
            ))

        vcsManager.setupSchedulers(self.config['schedulers'], spawner_name, try_name, deploy_name,
                                   self.importantManager, codebases, dep_properties)
        vcsManager.setupReporters(
            self.config['services'], spawner_name, try_name, codebases)
        res = vcsManager.setupChangeSource(self.config['services'])
        if res is not None:
            self.change_hook_dialects.update(res)
<commit_msg>Update configurator.py<commit_after>import urlparse
import os

from buildbot.config import error as config_error

# TBD use plugins!
from buildbot.config import BuilderConfig
from buildbot.schedulers.forcesched import StringParameter, CodebaseParameter
from buildbot.schedulers.triggerable import Triggerable
from buildbot.worker import Worker
from buildbot.worker import AbstractLatentWorker
from buildbot.process import factory
from buildbot.plugins import util
from buildbot import getVersion
from .important import ImportantManager
from .vcs import addRepository, getSupportedVCSTypes
from .steps import TravisSetupSteps
from .steps import TravisTrigger
from yaml import safe_load

import buildbot_travis

class TravisConfigurator(object):

    def __init__(self, config, vardir, latentRunners=False):
        self.config = config
        self.vardir = vardir
        self.latentRunners = latentRunners
        self.passwords = {}
        self.properties = {}
        self.repositories = {}
        self.cfgdict = {}
        self.importantManager = None
        self.change_hook_dialects = {}
        config.setdefault("builders", [])
        config.setdefault("schedulers", [])
        config.setdefault("change_source", [])
        config.setdefault("services", [])
        config.setdefault("status", [])
        self.defaultEnv = {}
        self.defaultStages = []
        # we are not really multimaster, but this remove some checks
        config.setdefault("multiMaster", True)
        config['codebaseGenerator'] = lambda chdict: chdict['project']
        self.config['title'] = os.environ.get(
            'buildbotTitle', "buildbot travis")

    def add_password(self, scheme, netloc, username, password):
        self.passwords[(scheme, netloc)] = (username, password)

    def fromYaml(self, path):
        buildbot_travis.api.setYamlPath(path)
        with open(path) as f:
            y = safe_load(f)

        return self.fromDict(y)

    def fromDict(self, y):
        buildbot_travis.api.setCfg(y)
        self.cfgdict = y
        self.importantManager = ImportantManager(
            y.get("not_important_files", []))
        self.defaultEnv = y.get("env", {})
        for k, v in self.defaultEnv.items():
            if not (isinstance(v, list) or isinstance(v, basestring)):
                config_error(
                    "'env' values must be strings or lists ; key %s is incorrect: %s" % (k, type(v)))
        for p in y.get("projects", []):
            self.define_travis_builder(**p)
        self.defaultStages = y.get("stages", [])
        for s in self.defaultStages:
            if not isinstance(s, basestring):
                config_error(
                    "'stages' values must be strings ; stage %s is incorrect: %s" % (s, type(s)))

        PORT = int(os.environ.get('PORT', 8020))
        self.config['buildbotURL'] = os.environ.get(
            'buildbotURL', "http://localhost:%d/" % (PORT, ))
        # minimalistic config to activate new web UI
        self.config['www'] = dict(port=PORT,
                                  change_hook_dialects=self.change_hook_dialects,
                                  plugins=dict(buildbot_travis={'cfg': self.cfgdict,
                                                                'supported_vcs': getSupportedVCSTypes()}),
                                  versions=[('buildbot_travis', getVersion(__file__))])

    def fromDb(self):
        buildbot_travis.api.useDbConfig()
        dbConfig = util.DbConfig(self.config, self.vardir)
        return self.fromDict(dbConfig.get("travis", {}))

    def get_spawner_workers(self):
        workers = [s.workername for s in self.config[
            'workers'] if isinstance(s, Worker)]
        return workers

    def get_runner_workers(self):
        if self.latentRunners:
            WorkerClass = AbstractLatentWorker
        else:
            WorkerClass = Worker
        workers = [s.workername for s in self.config[
            'workers'] if isinstance(s, WorkerClass)]
        return workers

    def define_travis_builder(self, name, repository, tags=None, **kwargs):
        name = str(name)
        repository = str(repository)
        job_name = "%s-job" % name
        try_name = "%s-try" % name
        deploy_name = "%s-deploy" % name
        spawner_name = name
        if tags is None:
            tags = []

        def formatTag(tag):
            if isinstance(tag, basestring):
                return str(tag)
            return str(tag['text'])

        tags = map(formatTag, tags)
        if 'username' not in kwargs and 'password' not in kwargs:
            p = urlparse.urlparse(repository)
            k = (p.scheme, p.netloc)
            if k in self.passwords:
                kwargs['username'], kwargs['password'] = self.passwords[k]

        codebases = {spawner_name: {'repository': repository}}
        for subrepo in kwargs.get('subrepos', []):
            codebases[subrepo['project']] = {
                'repository': subrepo['repository']}

        vcsManager = addRepository(
            name, dict(name=name, repository=repository, **kwargs))
        vcsManager.vardir = self.vardir

        # Define the builder for the main job
        f = factory.BuildFactory()
        vcsManager.addSourceSteps(f)
        f.addStep(TravisSetupSteps())

        self.config['builders'].append(BuilderConfig(
            name=job_name,
            workernames=self.get_runner_workers(),
            properties=self.properties,
            collapseRequests=False,
            env=self.defaultEnv,
            tags=["job", name] + tags,
            factory=f
        ))

        self.config['schedulers'].append(Triggerable(
            name=job_name,
            builderNames=[job_name],
            codebases=codebases,
        ))

        # Define the builder for a spawner
        f = factory.BuildFactory()
        vcsManager.addSourceSteps(f)
        f.addStep(TravisTrigger(
            scheduler=job_name,
        ))
        properties = dict(TRAVIS_PULL_REQUEST=False)
        properties.update(self.properties)
        self.config['builders'].append(BuilderConfig(
            name=spawner_name,
            workernames=self.get_spawner_workers(),
            properties=properties,
            tags=["trunk", name] + tags,
            factory=f
        ))

        # Define the builder for the deployment of the project
        f = factory.BuildFactory()
        vcsManager.addSourceSteps(f)
        f.addStep(TravisSetupSteps())

        # To manage deployment properly (with change traceability),
        # we need the version and the target deployment environment or "stage"
        version = StringParameter(name='version', label='GIT tag',
                                  hide=False, required=False, size=20)
        stage = StringParameter(name='stage', label='Stage',
                                hide=False, required=False,size=20)

        dep_properties = [version, stage]

        self.config['builders'].append(BuilderConfig(
            name=deploy_name,
            workernames=self.get_runner_workers(),
            tags=["deploy", name] + tags,
            factory=f
        ))

        if vcsManager.supportsTry:
            properties = dict(TRAVIS_PULL_REQUEST=True)
            properties.update(self.properties)
            # Define the builder for try job
            f = factory.BuildFactory()
            vcsManager.addSourceSteps(f)
            f.addStep(TravisTrigger(
                scheduler=job_name,
            ))

            self.config['builders'].append(BuilderConfig(
                name=try_name,
                workernames=self.get_spawner_workers(),
                properties=properties,
                tags=["try", name] + tags,
                factory=f
            ))

        vcsManager.setupSchedulers(self.config['schedulers'], spawner_name, try_name, deploy_name,
                                   self.importantManager, codebases, dep_properties)
        vcsManager.setupReporters(
            self.config['services'], spawner_name, try_name, codebases)
        res = vcsManager.setupChangeSource(self.config['services'])
        if res is not None:
            self.change_hook_dialects.update(res)
<|endoftext|>


<commit_before>from django.test import TestCase

from .models import OAuthScope
from .scoping import Scopes

class ScopingTestCase(TestCase):
    def setUp(self):
        self.s = Scopes()
        self.scope_a = OAuthScope.objects.create()
        self.scope_b = OAuthScope.objects.create()

    def test_add_scope(self):
        self.scope_a.scope_number = self.s.add_scope(
                                        self.scope_a.scope_number,
                                        "roombookings"
                                    )
        self.scope_a.scope_number = self.s.add_scope(
                                        self.scope_a.scope_number,
                                        "timetable"
                                    )
        self.scope_a.save()

        self.scope_b.scope_number = self.s.add_scope(
                                        self.scope_b.scope_number,
                                        "timetable"
                                    )
        self.scope_b.save()

        self.assertEqual(self.scope_a.scope_number, 3)
        self.assertEqual(self.scope_b.scope_number, 2)

    def test_remove_scope(self):
        self.scope_a.scope_number = 3
        self.scope_a.scope_number = self.s.remove_scope(
                                        self.scope_a.scope_number,
                                        "roombookings"
                                    )
        self.scope_a.save()

        self.assertEqual(self.scope_a.scope_number, 2)

    def test_scopes_equal(self):
        equal = self.scope_a.scopeIsEqual(self.scope_b)

        self.assertEqual(equal, True)

    def test_check_scope(self):
        self.assertEqual(self.s.check_scope(12, "roombookings"), False)
        self.assertEqual(self.s.check_scope(12, "uclu"), True)

    def test_get_all_scopes(self):
        scopes_bare = self.s.get_all_scopes(pretty_print=False)
        self.assertEqual({
            "name": "roombookings",
            "id": 0
        } in scopes_bare, True)

        scopes_pretty = self.s.get_all_scopes(pretty_print=True)
        self.assertEqual({
            "name": "roombookings",
            "description": "Private room bookings data"
        } in scopes_pretty, True)

    def test_scopes_dict(self):
        scopes_dict = self.s.scope_dict(5, pretty_print=False)
        self.assertEqual(scopes_dict, [
            {
                "name": "roombookings",
                "id": 0
            },
            {
                "name": "uclu",
                "id": 2
            }
        ])
<commit_msg>tests for oauth_token_check decorator (-permission)<commit_after>import json
import unittest.mock

from django.test import TestCase
from rest_framework.test import APIRequestFactory

from dashboard.models import App, User

from .decorators import oauth_token_check
from .models import OAuthScope, OAuthToken
from .scoping import Scopes

class ScopingTestCase(TestCase):
    def setUp(self):
        self.s = Scopes()
        self.scope_a = OAuthScope.objects.create()
        self.scope_b = OAuthScope.objects.create()

    def test_add_scope(self):
        self.scope_a.scope_number = self.s.add_scope(
                                        self.scope_a.scope_number,
                                        "roombookings"
                                    )
        self.scope_a.scope_number = self.s.add_scope(
                                        self.scope_a.scope_number,
                                        "timetable"
                                    )
        self.scope_a.save()

        self.scope_b.scope_number = self.s.add_scope(
                                        self.scope_b.scope_number,
                                        "timetable"
                                    )
        self.scope_b.save()

        self.assertEqual(self.scope_a.scope_number, 3)
        self.assertEqual(self.scope_b.scope_number, 2)

    def test_remove_scope(self):
        self.scope_a.scope_number = 3
        self.scope_a.scope_number = self.s.remove_scope(
                                        self.scope_a.scope_number,
                                        "roombookings"
                                    )
        self.scope_a.save()

        self.assertEqual(self.scope_a.scope_number, 2)

    def test_scopes_equal(self):
        equal = self.scope_a.scopeIsEqual(self.scope_b)

        self.assertEqual(equal, True)

    def test_check_scope(self):
        self.assertEqual(self.s.check_scope(12, "roombookings"), False)
        self.assertEqual(self.s.check_scope(12, "uclu"), True)

    def test_get_all_scopes(self):
        scopes_bare = self.s.get_all_scopes(pretty_print=False)
        self.assertEqual({
            "name": "roombookings",
            "id": 0
        } in scopes_bare, True)

        scopes_pretty = self.s.get_all_scopes(pretty_print=True)
        self.assertEqual({
            "name": "roombookings",
            "description": "Private room bookings data"
        } in scopes_pretty, True)

    def test_scopes_dict(self):
        scopes_dict = self.s.scope_dict(5, pretty_print=False)
        self.assertEqual(scopes_dict, [
            {
                "name": "roombookings",
                "id": 0
            },
            {
                "name": "uclu",
                "id": 2
            }
        ])

class OAuthTokenCheckDecoratorTestCase(TestCase):
    def setUp(self):
        mock_status_code = unittest.mock.Mock()
        mock_status_code.status_code = 200
        mock_view_func = unittest.mock.Mock(return_value=mock_status_code)

        self.dec_view = oauth_token_check(required_scopes=[])(mock_view_func)
        self.factory = APIRequestFactory()

    def test_decorator_no_client_secret_proof_provided(self):
        request = self.factory.get('/')
        response = self.dec_view(request)

        content = json.loads(response.content.decode())

        self.assertEqual(response.status_code, 400)
        self.assertFalse(content["ok"])
        self.assertEqual(content["error"], "No Client Secret Proof provided")

    def test_decorator_no_token_provided(self):
        request = self.factory.get(
            '/',
            {'client_secret_proof': 'not_a_real_proof'}
        )
        response = self.dec_view(request)

        content = json.loads(response.content.decode())

        self.assertEqual(response.status_code, 400)
        self.assertFalse(content["ok"])
        self.assertEqual(content["error"], "No token provided via GET.")

    def test_decorator_nonexistent_token_provided(self):
        request = self.factory.get(
            '/',
            {'client_secret_proof': 'not_a_real_secret_proof', 'token': 'fake'}
        )
        response = self.dec_view(request)

        content = json.loads(response.content.decode())

        self.assertEqual(response.status_code, 400)
        self.assertFalse(content["ok"])
        self.assertEqual(content["error"], "Token does not exist")

    def test_decorator_client_secret_verification_failed(self):
        # create User, App, and OAuth it
        user_ = User.objects.create(
            email="test@ucl.ac.uk", cn="test",
            given_name="Test Test"
        )
        app_ = App.objects.create(user=user_, name="An App")
        oauth_scope = OAuthScope.objects.create()
        oauth_token = OAuthToken.objects.create(
            app=app_, user=user_, scope=oauth_scope
        )

        request = self.factory.get(
            '/',
            {
                'client_secret_proof': 'not_a_real_proof',
                'token': oauth_token.token
            }
        )
        response = self.dec_view(request)

        content = json.loads(response.content.decode())

        self.assertEqual(response.status_code, 400)
        self.assertFalse(content["ok"])
        self.assertEqual(
            content["error"],
            "Client secret HMAC verification failed."
        )

    def test_decorator_no_permission_to_access(self):
        pass

    def test_decorator_everything_passes(self):
        user_ = User.objects.create(
            email="test@ucl.ac.uk", cn="test",
            given_name="Test Test"
        )
        app_ = App.objects.create(user=user_, name="An App")
        oauth_scope = OAuthScope.objects.create()
        oauth_token = OAuthToken.objects.create(
            app=app_, user=user_, scope=oauth_scope
        )

        # TODO: THIS WILL GET CHANGED
        # calculate the client_secret_proof
        import hmac, hashlib, base64
        hmac_digest = hmac.new(
            bytes(app_.client_secret, 'ascii'),
            msg=oauth_token.token.encode('ascii'),
            digestmod=hashlib.sha256
        ).digest()
        hmac_b64 = base64.b64encode(hmac_digest).decode()

        request = self.factory.get(
            '/',
            {'client_secret_proof': hmac_b64, 'token': oauth_token.token}
        )
        response = self.dec_view(request)

        self.assertEqual(response.status_code, 200)
<|endoftext|>


<commit_before>from django.shortcuts import render, redirect
from django.conf import settings
from django.core.files.storage import FileSystemStorage

from .models import Document
from .forms import DocumentForm

def index(request):
    return render(request,'jeopardyctf/index.html',{})

def upload(request):
    if request.method == 'POST' and request.FILES['myfile']:
        myfile = request.FILES['myfile']
        fs = FileSystemStorage()
        filename = fs.save(myfile.name, myfile)
        uploaded_file_url = fs.url(filename)
        return render(request, 'jeopardyctf/upload.html', {
            'uploaded_file_url': uploaded_file_url
        })
    return render(request, 'jeopardyctf/upload.html')

def upload2(request):
    if request.method == 'POST':
        form = DocumentForm(request.POST, request.FILES)
        if form.is_valid():
            form.save()
            return redirect('/jeopardy')
    else:
        form = DocumentForm()
    return render(request, 'jeopardyctf/upload2.html', {
        'form': form
    })
<commit_msg>Listing Uploaded Files<commit_after>from django.shortcuts import render, redirect
from django.conf import settings
from django.core.files.storage import FileSystemStorage

from .models import Document
from .forms import DocumentForm

def index(request):
    path=settings.MEDIA_ROOT
    file_list =os.listdir(path)
    return render_to_response('index.html', {'files': file_list})

def upload(request):
    if request.method == 'POST' and request.FILES['myfile']:
        myfile = request.FILES['myfile']
        fs = FileSystemStorage()
        filename = fs.save(myfile.name, myfile)
        uploaded_file_url = fs.url(filename)
        return render(request, 'jeopardyctf/upload.html', {
            'uploaded_file_url': uploaded_file_url
        })
    return render(request, 'jeopardyctf/upload.html')

def upload2(request):
    if request.method == 'POST':
        form = DocumentForm(request.POST, request.FILES)
        if form.is_valid():
            form.save()
            return redirect('/jeopardy')
    else:
        form = DocumentForm()
    return render(request, 'jeopardyctf/upload2.html', {
        'form': form
    })
<|endoftext|>


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

import xml.etree.cElementTree as ET
import re

class OSMDataDescriptor(object):
    '''
    OSMDataDescriptor.
    Simple class to summarize OSM data.
    '''

    def __init__(self, filename):
        self.filename = filename

    def count_element(self):
        '''Return element and its total.'''
        tags = {}
        context = ET.iterparse(self.filename)
        for event, element in iter(context):
            tags[element.tag] = tags.get(element.tag, 0) + 1

        return tags

    def get_tag_keys(self):
        '''Return list of sorted tag of key attribute and its total.'''
        tag_keys = {}
        for event, elem in ET.iterparse(self.filename, events=('start',)):
            for tag in elem.iter("tag"):
                tag_keys[tag.attrib['k']] = tag_keys.get(tag.attrib['k'], 0) + 1

        return sorted(tag_keys.items(), key=lambda x: x[0])

    def audit_abbreviated_street_prefixes(self):
        '''Check the address prefixes for street.

        Faulty street address will startswith:
            jln, jln., jl, jl.
        '''

        street_regx = re.compile(r'^(jln?\.?)\s?(.*)', re.IGNORECASE)
        faulty_street_prefixes = []

        for event, element in ET.iterparse(self.filename, events=('start',)):
            # Only pay attention to 'way' tag
            if not element.tag == 'way':
                continue

            for el in element.iter('tag'):
                # We audit both addr:street and addr:full, since
                # a lot of address is in the addr:full tag.
                if el.attrib['k'] == 'addr:street' or el.attrib['k'] == 'addr:full':
                    if street_regx.match(el.attrib['v']):
                        faulty_street_prefixes.append(el.attrib['v'])

        return faulty_street_prefixes

    def audit_abbreviated_alley_prefixes(self):
        '''Check the alley prefixes for street.

        Faulty alley address will starts with:
            gg, gg.
        '''

        alley_regx = re.compile(r'^(gg\.?)\s?(.*)', re.IGNORECASE)
        faulty_alley_prefixes = []

        for event, element in ET.iterparse(self.filename, events=('start',)):
            # Only pay attention to 'way' tag
            if not element.tag == 'way':
                continue

            for el in element.iter('tag'):
                # We audit both addr:street and addr:full, since
                # a lot of address is in the addr:full tag.
                if el.attrib['k'] == 'addr:street' or el.attrib['k'] == 'addr:full':
                    if alley_regx.match(el.attrib['v']):
                        faulty_alley_prefixes.append(el.attrib['v'])

        return faulty_alley_prefixes
<commit_msg>Add similiar address name.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import xml.etree.cElementTree as ET
import re
import json
import itertools
from difflib import SequenceMatcher

class OSMDataDescriptor(object):
    '''
    OSMDataDescriptor.
    Simple class to summarize OSM data.
    '''

    def __init__(self, filename):
        self.filename = filename

    def count_element(self):
        '''Return element and its total.'''
        tags = {}
        context = ET.iterparse(self.filename)
        for event, element in iter(context):
            tags[element.tag] = tags.get(element.tag, 0) + 1

        return tags

    def get_tag_keys(self):
        '''Return list of sorted tag of key attribute and its total.'''
        tag_keys = {}
        for event, elem in ET.iterparse(self.filename, events=('start',)):
            for tag in elem.iter("tag"):
                tag_keys[tag.attrib['k']] = tag_keys.get(tag.attrib['k'], 0) + 1

        return sorted(tag_keys.items(), key=lambda x: x[0])

    def audit_abbreviated_street_prefixes(self):
        '''Check the address prefixes for street.

        Faulty street address will startswith:
            jln, jln., jl, jl.
        '''

        street_regx = re.compile(r'^(jln?\.?)\s?(.*)', re.IGNORECASE)
        faulty_street_prefixes = []

        for event, element in ET.iterparse(self.filename, events=('start',)):
            # Only pay attention to 'way' tag
            if not element.tag == 'way':
                continue

            for el in element.iter('tag'):
                # We audit both addr:street and addr:full, since
                # a lot of address is in the addr:full tag.
                if el.attrib['k'] == 'addr:street' or el.attrib['k'] == 'addr:full':
                    if street_regx.match(el.attrib['v']):
                        faulty_street_prefixes.append(el.attrib['v'])

        return faulty_street_prefixes

    def audit_abbreviated_alley_prefixes(self):
        '''Check the alley prefixes for street.

        Faulty alley address will starts with:
            gg, gg.
        '''

        alley_regx = re.compile(r'^(gg\.?)\s?(.*)', re.IGNORECASE)
        faulty_alley_prefixes = []

        for event, element in ET.iterparse(self.filename, events=('start',)):
            # Only pay attention to 'way' tag
            if not element.tag == 'way':
                continue

            for el in element.iter('tag'):
                # We audit both addr:street and addr:full, since
                # a lot of address is in the addr:full tag.
                if el.attrib['k'] == 'addr:street' or el.attrib['k'] == 'addr:full':
                    if alley_regx.match(el.attrib['v']):
                        faulty_alley_prefixes.append(el.attrib['v'])

        return faulty_alley_prefixes

    def audit_address_similar_names(self, reference_filename='street_name.json'):
        '''Check if the street name match reference from wikipedia.

        The reference is taken from Wikipedia info box.
        Use difflib.SequenceMatcher to check how closely the name in the tag match the reference.
        '''
        # Load wikipedia street name data
        street_names_file = open(reference_filename)
        street_names = json.load(street_names_file)

        # Flatten the street names, ignore the city
        flatten_street_names = [item.lower() for item in itertools.chain(*street_names.values())]

        closely_matched_names = set()
        for event, element in ET.iterparse(self.filename, events=('start',)):
            # Only pay attention to 'way' tag
            if not element.tag == 'way':
                continue

            for el in element.iter('tag'):
                # We only audit addr:street for this.
                if el.attrib['k'] == 'addr:street':
                    street_name = el.attrib['v'].lower()
                    for ref_street_name in flatten_street_names:
                        s = SequenceMatcher(lambda x: x == " ", street_name, ref_street_name)

                        # Check the ratio between the two stings.
                        # 0.6 is considered a close match,
                        # 1 means the two strings match exactly.
                        # We want ratio that is between the two.
                        if 0.65 <= s.ratio() < 1.0:
                            closely_matched_names.add((ref_street_name, street_name, s.ratio()))

        return closely_matched_names
<|endoftext|>


<commit_before>
class Player:
    VERSION = "We're kicking your ass!"
    name = "weneverwonanything"
    stack = 0
    current_bid = 0
    player = True
    bet = 0
    hand = True
    good_cards = ["10", "J", "Q", "K", "A"]
    counter = 0
    common_cards = []
    suit = False

    def get_info(self, gs):
        self.current_bid = gs.get("current_buy_in")
        players = gs.get("players")
        self.player = players[gs.get("in_action")]
        self.stack = self.player.get("stack")
        self.bet = self.player.get("bet")
        self.hand = self.player.get("hole_cards")
        self.common_cards = gs.get("community_cards")
        self.suit = self.hand[0]["suit"] == self.hand[1]["suit"]
        pass

    def betRequest(self, gs):
        self.get_info(gs)
        while(self.counter == 0):
            if self.hand[0]["rank"] in self.good_cards:
                if self.hand[0]["rank"] == self.hand[1]["rank"]:
                    self.counter += 1
                    return 200
            if self.hand[1]["rank"] in self.good_cards:
                return 10
            return 0
        if len(self.common_cards) == 3:
            if (self.suit) and self.hand[0]["suit"] in self.common_cards:
                return 200
        if self.hand[0]["rank"] in self.good_cards and self.hand[1]["rank"] in self.good_cards:
            if len(self.common_cards) == 0:
                return 100
            if len(self.common_cards) >= 3:
                if self.hand[0]["rank"] in self.common_cards.values() and self.hand[1]["rank"] in self.common_cards.values():
                    return 10000
                if self.hand[0]["rank"] in self.common_cards.values() or self.hand[1]["rank"] in self.common_cards.values():
                    if self.hand[0]["rank"] == self.hand[1]["rank"]:
                        return 10000
                    for i in range(len(self.common_cards)) :
                        for j in range(len(self.common_cards) - 1):
                            if i != j+1:
                                if self.common_cards[i].get("rank") == self.common_cards[j+1].get("rank"):
                                    return 10000
                    return 300
            return 0
        return 0

    def showdown(self, game_state):
        pass
<commit_msg>adjust numbers<commit_after>
class Player:
    VERSION = "We're kicking your ass!"
    name = "weneverwonanything"
    stack = 0
    current_bid = 0
    player = True
    bet = 0
    hand = True
    good_cards = ["10", "J", "Q", "K", "A"]
    counter = 0
    common_cards = []
    suit = False

    def get_info(self, gs):
        self.current_bid = gs.get("current_buy_in")
        players = gs.get("players")
        self.player = players[gs.get("in_action")]
        self.stack = self.player.get("stack")
        self.bet = self.player.get("bet")
        self.hand = self.player.get("hole_cards")
        self.common_cards = gs.get("community_cards")
        self.suit = self.hand[0]["suit"] == self.hand[1]["suit"]
        pass

    def betRequest(self, gs):
        self.get_info(gs)
        while(self.counter == 0):
            if self.hand[0]["rank"] in self.good_cards:
                if self.hand[0]["rank"] == self.hand[1]["rank"]:
                    self.counter += 1
                    return 200
            if self.hand[1]["rank"] in self.good_cards:
                return 10
            return 0
        if len(self.common_cards) == 3:
            if (self.suit) and self.hand[0]["suit"] in self.common_cards:
                return 200
        if self.hand[0]["rank"] in self.good_cards and self.hand[1]["rank"] in self.good_cards:
            if len(self.common_cards) == 0:
                return 100
            if len(self.common_cards) >= 3:
                if self.hand[0]["rank"] in self.common_cards.values() and self.hand[1]["rank"] in self.common_cards.values():
                    return 300
                if self.hand[0]["rank"] in self.common_cards.values() or self.hand[1]["rank"] in self.common_cards.values():
                    if self.hand[0]["rank"] == self.hand[1]["rank"]:
                        return 10000
                    for i in range(len(self.common_cards)) :
                        for j in range(len(self.common_cards) - 1):
                            if i != j+1:
                                if self.common_cards[i].get("rank") == self.common_cards[j+1].get("rank"):
                                    return 10000
                    return 300
            return 0
        return 0

    def showdown(self, game_state):
        pass
<|endoftext|>


<commit_before>import pygame
from pygame.locals import *
from Box2D import *

from sprite import Sprite, Animation
from debug import draw_body

class Player(object):
    def __init__(self, world, pos=(0, 0)):
        self.body = world.b2world.CreateDynamicBody(
            position=pos,
            allowSleep=False
        )
        self.body.CreatePolygonFixture(
            box=(0.8, 1.6),
            density=1,
            friction=0.0,
        )
        self.body.linearDamping = 0.0
        self.body.angularDamping = 0.0
        self.sensor = self.body.CreatePolygonFixture(
            box=(0.6, 0.2, (0, -1.6), 0),
            density=1,
            friction=0.0,
            isSensor = True
        )

        self.stand_sprite = Sprite("data/art/stand.png")
        self.jump_sprite = Sprite("data/art/jump.png")
        self.fall_sprite = Sprite("data/art/fall.png")
        walksprite0 = self.stand_sprite
        walksprite1 = Sprite("data/art/walk1.png")
        walksprite2 = Sprite("data/art/walk2.png")
        self.walk_sprite = Animation((
            walksprite0, walksprite1, walksprite1,
            walksprite0, walksprite2, walksprite2
        ))
        self.sprite = self.stand_sprite

    @property
    def pos(self):
        return tuple(self.body.position)

    def update(self, delta):
        keys = pygame.key.get_pressed()

        can_jump = False
        for ce in self.body.contacts:
            if ce.contact.touching:
                fixtureA = ce.contact.fixtureA
                fixtureB = ce.contact.fixtureB
                if self.sensor == fixtureA or self.sensor == fixtureB:
                    can_jump = True
                    break

        # Vertical logic
        if keys[K_w] and can_jump:
            self.body.linearVelocity[1] = 30.0
        elif self.body.linearVelocity[1] > 0.0:
            vel_delta = self.body.linearVelocity[1] * delta * 2.0
            self.body.linearVelocity[1] -= vel_delta

        # Horizontal logic
        if keys[K_a]:
            self.body.linearVelocity[0] -= 30.0 * delta
        if keys[K_d]:
            self.body.linearVelocity[0] += 30.0 * delta
        self.body.linearVelocity[0] -= self.body.linearVelocity[0] * delta * 2.0

        self.body.angle = 0.0
        self.body.angularVelocity = 0.0

        # Sprite logic
        if can_jump:
            if abs(self.body.linearVelocity[0]) > 1.0:
                self.sprite = self.walk_sprite
            else:
                self.sprite = self.stand_sprite
        elif self.body.linearVelocity[1] > 0.0:
            self.sprite = self.jump_sprite
        else:
            self.sprite = self.fall_sprite
        sprite_delta = delta * abs(self.body.linearVelocity[0]) * 2.0
        mirror = self.body.linearVelocity[0] < 0.0
        self.sprite.update(sprite_delta, mirror)

    def render(self, screen, camera):
        self.sprite.render(screen, camera, self.pos)
        #draw_body(self.body, screen, camera)
<commit_msg>Add arrow keys to player movement (oops)<commit_after>import pygame
from pygame.locals import *
from Box2D import *

from sprite import Sprite, Animation
from debug import draw_body

class Player(object):
    def __init__(self, world, pos=(0, 0)):
        self.body = world.b2world.CreateDynamicBody(
            position=pos,
            allowSleep=False
        )
        self.body.CreatePolygonFixture(
            box=(0.8, 1.6),
            density=1,
            friction=0.0,
        )
        self.body.linearDamping = 0.0
        self.body.angularDamping = 0.0
        self.sensor = self.body.CreatePolygonFixture(
            box=(0.6, 0.2, (0, -1.6), 0),
            density=1,
            friction=0.0,
            isSensor = True
        )

        self.stand_sprite = Sprite("data/art/stand.png")
        self.jump_sprite = Sprite("data/art/jump.png")
        self.fall_sprite = Sprite("data/art/fall.png")
        walksprite0 = self.stand_sprite
        walksprite1 = Sprite("data/art/walk1.png")
        walksprite2 = Sprite("data/art/walk2.png")
        self.walk_sprite = Animation((
            walksprite0, walksprite1, walksprite1,
            walksprite0, walksprite2, walksprite2
        ))
        self.sprite = self.stand_sprite

    @property
    def pos(self):
        return tuple(self.body.position)

    def update(self, delta):
        keys = pygame.key.get_pressed()

        can_jump = False
        for ce in self.body.contacts:
            if ce.contact.touching:
                fixtureA = ce.contact.fixtureA
                fixtureB = ce.contact.fixtureB
                if self.sensor == fixtureA or self.sensor == fixtureB:
                    can_jump = True
                    break

        # Vertical logic
        if (keys[K_w] or keys[K_UP]) and can_jump:
            self.body.linearVelocity[1] = 30.0
        elif self.body.linearVelocity[1] > 0.0:
            vel_delta = self.body.linearVelocity[1] * delta * 2.0
            self.body.linearVelocity[1] -= vel_delta

        # Horizontal logic
        if keys[K_a] or keys[K_LEFT]:
            self.body.linearVelocity[0] -= 30.0 * delta
        if keys[K_d] or keys[K_RIGHT]:
            self.body.linearVelocity[0] += 30.0 * delta
        self.body.linearVelocity[0] -= self.body.linearVelocity[0] * delta * 2.0

        self.body.angle = 0.0
        self.body.angularVelocity = 0.0

        # Sprite logic
        if can_jump:
            if abs(self.body.linearVelocity[0]) > 1.0:
                self.sprite = self.walk_sprite
            else:
                self.sprite = self.stand_sprite
        elif self.body.linearVelocity[1] > 0.0:
            self.sprite = self.jump_sprite
        else:
            self.sprite = self.fall_sprite
        sprite_delta = delta * abs(self.body.linearVelocity[0]) * 2.0
        mirror = self.body.linearVelocity[0] < 0.0
        self.sprite.update(sprite_delta, mirror)

    def render(self, screen, camera):
        self.sprite.render(screen, camera, self.pos)
        #draw_body(self.body, screen, camera)
<|endoftext|>