<commit_before>"""Utilities to manipulate JSON objects.
"""
#-----------------------------------------------------------------------------
#  Copyright (C) 2010-2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING.txt, distributed as part of this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------
# stdlib
import re
import sys
import types
from base64 import encodestring
from datetime import datetime

from IPython.utils import py3compat
from IPython.utils.encoding import DEFAULT_ENCODING
from IPython.utils import text
next_attr_name = '__next__' if py3compat.PY3 else 'next'

#-----------------------------------------------------------------------------
# Globals and constants
#-----------------------------------------------------------------------------

# timestamp formats
ISO8601="%Y-%m-%dT%H:%M:%S.%f"
ISO8601_PAT=re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d+$")

#-----------------------------------------------------------------------------
# Classes and functions
#-----------------------------------------------------------------------------

def rekey(dikt):
    """Rekey a dict that has been forced to use str keys where there should be
    ints by json."""
    for k in dikt.iterkeys():
        if isinstance(k, basestring):
            ik=fk=None
            try:
                ik = int(k)
            except ValueError:
                try:
                    fk = float(k)
                except ValueError:
                    continue
            if ik is not None:
                nk = ik
            else:
                nk = fk
            if nk in dikt:
                raise KeyError("already have key %r"%nk)
            dikt[nk] = dikt.pop(k)
    return dikt

def extract_dates(obj):
    """extract ISO8601 dates from unpacked JSON"""
    if isinstance(obj, dict):
        obj = dict(obj) # don't clobber
        for k,v in obj.iteritems():
            obj[k] = extract_dates(v)
    elif isinstance(obj, (list, tuple)):
        obj = [ extract_dates(o) for o in obj ]
    elif isinstance(obj, basestring):
        if ISO8601_PAT.match(obj):
            obj = datetime.strptime(obj, ISO8601)
    return obj

def squash_dates(obj):
    """squash datetime objects into ISO8601 strings"""
    if isinstance(obj, dict):
        obj = dict(obj) # don't clobber
        for k,v in obj.iteritems():
            obj[k] = squash_dates(v)
    elif isinstance(obj, (list, tuple)):
        obj = [ squash_dates(o) for o in obj ]
    elif isinstance(obj, datetime):
        obj = obj.strftime(ISO8601)
    return obj

def date_default(obj):
    """default function for packing datetime objects in JSON."""
    if isinstance(obj, datetime):
        return obj.strftime(ISO8601)
    else:
        raise TypeError("%r is not JSON serializable"%obj)

# constants for identifying png/jpeg data
PNG = b'\x89PNG\r\n\x1a\n'
JPEG = b'\xff\xd8'

def encode_images(format_dict):
    """b64-encodes images in a displaypub format dict

    Perhaps this should be handled in json_clean itself?
    """
    encoded = format_dict.copy()
    pngdata = format_dict.get('image/png')
    if isinstance(pngdata, bytes) and pngdata[:8] == PNG:
        encoded['image/png'] = encodestring(pngdata).decode('ascii')
    jpegdata = format_dict.get('image/jpeg')
    if isinstance(jpegdata, bytes) and jpegdata[:2] == JPEG:
        encoded['image/jpeg'] = encodestring(jpegdata).decode('ascii')
    return encoded

def json_clean(obj):
    """Clean an object to ensure it's safe to encode in JSON.

    Atomic, immutable objects are returned unmodified.  Sets and tuples are
    converted to lists, lists are copied and dicts are also copied.

    Note: dicts whose keys could cause collisions upon encoding (such as a dict
    with both the number 1 and the string '1' as keys) will cause a ValueError
    to be raised.

    Parameters
    ----------
    obj : any python object

    Returns
    -------
    out : object

      A version of the input which will not cause an encoding error when
      encoded as JSON.  Note that this function does not *encode* its inputs,
      it simply sanitizes it so that there will be no encoding errors later.

    Examples
    --------
    >>> json_clean(4)
    4
    >>> json_clean(range(10))
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    >>> sorted(json_clean(dict(x=1, y=2)).items())
    [('x', 1), ('y', 2)]
    >>> sorted(json_clean(dict(x=1, y=2, z=[1,2,3])).items())
    [('x', 1), ('y', 2), ('z', [1, 2, 3])]
    >>> json_clean(True)
    True
    """
    # types that are 'atomic' and ok in json as-is.  bool doesn't need to be
    # listed explicitly because bools pass as int instances
    atomic_ok = (unicode, int, float, types.NoneType)

    # containers that we need to convert into lists
    container_to_list = (tuple, set, types.GeneratorType)

    if isinstance(obj, atomic_ok):
        return obj

    if isinstance(obj, bytes):
        return obj.decode(DEFAULT_ENCODING, 'replace')

    if isinstance(obj, container_to_list) or (
        hasattr(obj, '__iter__') and hasattr(obj, next_attr_name)):
        obj = list(obj)

    if isinstance(obj, list):
        return [json_clean(x) for x in obj]

    if isinstance(obj, dict):
        # First, validate that the dict won't lose data in conversion due to
        # key collisions after stringification.  This can happen with keys like
        # True and 'true' or 1 and '1', which collide in JSON.
        nkeys = len(obj)
        nkeys_collapsed = len(set(map(str, obj)))
        if nkeys != nkeys_collapsed:
            raise ValueError('dict can not be safely converted to JSON: '
                             'key collision would lead to dropped values')
        # If all OK, proceed by making the new dict that will be json-safe
        out = {}
        for k,v in obj.iteritems():
            out[str(k)] = json_clean(v)
        return out

    # If we get here, we don't know how to handle the object, so we just get
    # its repr and return that.  This will catch lambdas, open sockets, class
    # objects, and any other complicated contraption that json can't encode
    return repr(obj)
<commit_msg>expand encode_images docstring<commit_after>"""Utilities to manipulate JSON objects.
"""
#-----------------------------------------------------------------------------
#  Copyright (C) 2010-2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING.txt, distributed as part of this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------
# stdlib
import re
import sys
import types
from base64 import encodestring
from datetime import datetime

from IPython.utils import py3compat
from IPython.utils.encoding import DEFAULT_ENCODING
from IPython.utils import text
next_attr_name = '__next__' if py3compat.PY3 else 'next'

#-----------------------------------------------------------------------------
# Globals and constants
#-----------------------------------------------------------------------------

# timestamp formats
ISO8601="%Y-%m-%dT%H:%M:%S.%f"
ISO8601_PAT=re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d+$")

#-----------------------------------------------------------------------------
# Classes and functions
#-----------------------------------------------------------------------------

def rekey(dikt):
    """Rekey a dict that has been forced to use str keys where there should be
    ints by json."""
    for k in dikt.iterkeys():
        if isinstance(k, basestring):
            ik=fk=None
            try:
                ik = int(k)
            except ValueError:
                try:
                    fk = float(k)
                except ValueError:
                    continue
            if ik is not None:
                nk = ik
            else:
                nk = fk
            if nk in dikt:
                raise KeyError("already have key %r"%nk)
            dikt[nk] = dikt.pop(k)
    return dikt

def extract_dates(obj):
    """extract ISO8601 dates from unpacked JSON"""
    if isinstance(obj, dict):
        obj = dict(obj) # don't clobber
        for k,v in obj.iteritems():
            obj[k] = extract_dates(v)
    elif isinstance(obj, (list, tuple)):
        obj = [ extract_dates(o) for o in obj ]
    elif isinstance(obj, basestring):
        if ISO8601_PAT.match(obj):
            obj = datetime.strptime(obj, ISO8601)
    return obj

def squash_dates(obj):
    """squash datetime objects into ISO8601 strings"""
    if isinstance(obj, dict):
        obj = dict(obj) # don't clobber
        for k,v in obj.iteritems():
            obj[k] = squash_dates(v)
    elif isinstance(obj, (list, tuple)):
        obj = [ squash_dates(o) for o in obj ]
    elif isinstance(obj, datetime):
        obj = obj.strftime(ISO8601)
    return obj

def date_default(obj):
    """default function for packing datetime objects in JSON."""
    if isinstance(obj, datetime):
        return obj.strftime(ISO8601)
    else:
        raise TypeError("%r is not JSON serializable"%obj)

# constants for identifying png/jpeg data
PNG = b'\x89PNG\r\n\x1a\n'
JPEG = b'\xff\xd8'

def encode_images(format_dict):
    """b64-encodes images in a displaypub format dict

    Perhaps this should be handled in json_clean itself?

    Parameters
    ----------

    format_dict : dict
        A dictionary of display data keyed by mime-type

    Returns
    -------

    format_dict : dict
        A copy of the same dictionary,
        but binary image data ('image/png' or 'image/jpeg')
        is base64-encoded.

    """
    encoded = format_dict.copy()
    pngdata = format_dict.get('image/png')
    if isinstance(pngdata, bytes) and pngdata[:8] == PNG:
        encoded['image/png'] = encodestring(pngdata).decode('ascii')
    jpegdata = format_dict.get('image/jpeg')
    if isinstance(jpegdata, bytes) and jpegdata[:2] == JPEG:
        encoded['image/jpeg'] = encodestring(jpegdata).decode('ascii')
    return encoded

def json_clean(obj):
    """Clean an object to ensure it's safe to encode in JSON.

    Atomic, immutable objects are returned unmodified.  Sets and tuples are
    converted to lists, lists are copied and dicts are also copied.

    Note: dicts whose keys could cause collisions upon encoding (such as a dict
    with both the number 1 and the string '1' as keys) will cause a ValueError
    to be raised.

    Parameters
    ----------
    obj : any python object

    Returns
    -------
    out : object

      A version of the input which will not cause an encoding error when
      encoded as JSON.  Note that this function does not *encode* its inputs,
      it simply sanitizes it so that there will be no encoding errors later.

    Examples
    --------
    >>> json_clean(4)
    4
    >>> json_clean(range(10))
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    >>> sorted(json_clean(dict(x=1, y=2)).items())
    [('x', 1), ('y', 2)]
    >>> sorted(json_clean(dict(x=1, y=2, z=[1,2,3])).items())
    [('x', 1), ('y', 2), ('z', [1, 2, 3])]
    >>> json_clean(True)
    True
    """
    # types that are 'atomic' and ok in json as-is.  bool doesn't need to be
    # listed explicitly because bools pass as int instances
    atomic_ok = (unicode, int, float, types.NoneType)

    # containers that we need to convert into lists
    container_to_list = (tuple, set, types.GeneratorType)

    if isinstance(obj, atomic_ok):
        return obj

    if isinstance(obj, bytes):
        return obj.decode(DEFAULT_ENCODING, 'replace')

    if isinstance(obj, container_to_list) or (
        hasattr(obj, '__iter__') and hasattr(obj, next_attr_name)):
        obj = list(obj)

    if isinstance(obj, list):
        return [json_clean(x) for x in obj]

    if isinstance(obj, dict):
        # First, validate that the dict won't lose data in conversion due to
        # key collisions after stringification.  This can happen with keys like
        # True and 'true' or 1 and '1', which collide in JSON.
        nkeys = len(obj)
        nkeys_collapsed = len(set(map(str, obj)))
        if nkeys != nkeys_collapsed:
            raise ValueError('dict can not be safely converted to JSON: '
                             'key collision would lead to dropped values')
        # If all OK, proceed by making the new dict that will be json-safe
        out = {}
        for k,v in obj.iteritems():
            out[str(k)] = json_clean(v)
        return out

    # If we get here, we don't know how to handle the object, so we just get
    # its repr and return that.  This will catch lambdas, open sockets, class
    # objects, and any other complicated contraption that json can't encode
    return repr(obj)
<|endoftext|>


<commit_before>import sublime, sublime_plugin
import os, zipfile

class Schemr():
	def load_schemes(self):
		color_schemes = []

		for root, dirs, files in os.walk(sublime.packages_path()):
			for filename in (filename for filename in files if filename.endswith('.tmTheme')):
				name = filename.replace('.tmTheme', '')
				filepath = os.path.join(root, filename).replace(sublime.packages_path(), 'Packages').replace('\\', '/')
				color_schemes.append([name, filepath])

		for root, dirs, files in os.walk(sublime.installed_packages_path()):
			for filename in (filename for filename in files if filename.startswith('Color Scheme - ') and filename.endswith('.sublime-package')):
				package = zipfile.ZipFile(os.path.join(sublime.installed_packages_path(), filename))
				for f in (f for f in package.namelist() if f.endswith('.tmTheme')):
					name = f.replace('.tmTheme', '')
					filepath = os.path.join(root, filename, f).replace(sublime.installed_packages_path(), 'Packages').replace('.sublime-package', '').replace('\\', '/')
					color_schemes.append([name, filepath])

		return color_schemes

	def set_scheme(self, s):
		self.settings().set('color_scheme', s)
		sublime.save_settings('Preferences.sublime-settings')

	def get_scheme(self):
		return self.settings().get('color_scheme')

	def cycle_scheme(self, d):
		color_schemes = self.load_schemes()
		the_scheme = self.get_scheme()
		the_index = [scheme[1] for scheme in color_schemes].index(the_scheme)
		num_of_schemes = len(color_schemes)

		if d == 1:
			index = the_index + 1 if the_index < num_of_schemes - 1 else 0

		if d == -1:
			index = the_index - 1 if the_index > 0 else num_of_schemes - 1

		self.set_scheme(color_schemes[index][1])
		sublime.status_message('Schemr: ' + color_schemes[index][0])

	def settings(self):
		return sublime.load_settings('Preferences.sublime-settings')

Schemr = Schemr()

class SchemrListSchemesCommand(sublime_plugin.WindowCommand):
	def run(self):
		color_schemes = Schemr.load_schemes()

		def on_done(index):
			if index != -1:
				Schemr.set_scheme(color_schemes[index][1])
				sublime.status_message('Schemr: ' + color_schemes[index][0])

		self.window.show_quick_panel(color_schemes, on_done)

class SchemrNextSchemeCommand(sublime_plugin.WindowCommand):
	def run(self):
		Schemr.cycle_scheme(1)

class SchemrPrevSchemeCommand(sublime_plugin.WindowCommand):
	def run(self):
		Schemr.cycle_scheme(-1)
<commit_msg>Rename vars for readability<commit_after>import sublime, sublime_plugin
import os, zipfile

class Schemr():
	def load_schemes(self):
		color_schemes = []

		for root, dirs, files in os.walk(sublime.packages_path()):
			for filename in (filename for filename in files if filename.endswith('.tmTheme')):
				name = filename.replace('.tmTheme', '')
				filepath = os.path.join(root, filename).replace(sublime.packages_path(), 'Packages').replace('\\', '/')
				color_schemes.append(['Scheme: ' + name, filepath])

		for root, dirs, files in os.walk(sublime.installed_packages_path()):
			for package in (package for package in files if package.endswith('.sublime-package')):
				zf = zipfile.ZipFile(os.path.join(sublime.installed_packages_path(), package))
				for filename in (filename for filename in zf.namelist() if filename.endswith('.tmTheme')):
					name = os.path.basename(filename).replace('.tmTheme', '')
					filepath = os.path.join(root, package, filename).replace(sublime.installed_packages_path(), 'Packages').replace('.sublime-package', '').replace('\\', '/')
					color_schemes.append(['Scheme: ' + name, filepath])

		return color_schemes

	def set_scheme(self, s):
		self.settings().set('color_scheme', s)
		sublime.save_settings('Preferences.sublime-settings')

	def get_scheme(self):
		return self.settings().get('color_scheme')

	def cycle_scheme(self, d):
		color_schemes = self.load_schemes()
		the_scheme = self.get_scheme()
		the_index = [scheme[1] for scheme in color_schemes].index(the_scheme)
		num_of_schemes = len(color_schemes)

		if d == 1:
			index = the_index + 1 if the_index < num_of_schemes - 1 else 0

		if d == -1:
			index = the_index - 1 if the_index > 0 else num_of_schemes - 1

		self.set_scheme(color_schemes[index][1])
		sublime.status_message(color_schemes[index][0])

	def settings(self):
		return sublime.load_settings('Preferences.sublime-settings')

Schemr = Schemr()

class SchemrListSchemesCommand(sublime_plugin.WindowCommand):
	def run(self):
		color_schemes = Schemr.load_schemes()

		def on_done(index):
			if index != -1:
				Schemr.set_scheme(color_schemes[index][1])
				sublime.status_message(color_schemes[index][0])

		self.window.show_quick_panel(color_schemes, on_done)

class SchemrNextSchemeCommand(sublime_plugin.WindowCommand):
	def run(self):
		Schemr.cycle_scheme(1)

class SchemrPrevSchemeCommand(sublime_plugin.WindowCommand):
	def run(self):
		Schemr.cycle_scheme(-1)
<|endoftext|>


<commit_before>from src.websites.kickassanime import Scraper
from src.utils.webdriver import get_chrome_webdriver
from src.scrape_utils.servers import StreamServers

url = "https://www.kickassanime.ru/anime/clannad"
episodes = []
driver = get_chrome_webdriver()
scraper = Scraper(driver, url, StreamServers.YOURUPLOAD)

fetched_episodes = {}

meta_data = scraper.fetch_metadata() # -> { title: _____, synopsis: _____, poster: ____ }

# Parse user input
# Options: Anime URL, Start episode, End episode, Missing -> [Use metadata, Use file list], Auto-add, Metadata only
# Read config
# Check URL and get matching scraper instance
# Fetch episodes from prefered server
# Prepare JSON and simple text files

'''
if episodes:
    for ep in episodes: # TODO: fix this
        fetched_episodes += [scraper.fetch_episode(ep)] # -> { stream_page: ____, stream_url: ___ }
else:
    scraper.fetch_all_episodes(fetched_episodes)

driver.close()

print(fetched_episodes)

# New Anime object
# Serialize object
# Output operable data

'''<commit_msg>Basic CLI (barely working)<commit_after>import sys

from argparse import ArgumentParser

from src.websites.kickassanime import Scraper
from src.utils.webdriver import get_chrome_webdriver
from src.scrape_utils.servers import StreamServers

parser = ArgumentParser()
parser.add_argument("url", help="URL to the page of the list of episodes of the anime")

parser.add_argument("--start", "-s", type=int, help="The episode number to start fetching from")
parser.add_argument("--end", "-e", type=int, help="The episode number to stop fetching at")
parser.add_argument("--missing", "-m", help="Fetch downloads URLs only for episodes not present in this directory",
                        action="store_true")

parser.add_argument("--auto", "-a", help="Automatically add the downloads to IDM using the current directory as the download location",
                        action="store_true")

args = parser.parse_args()
url = args.url
is_auto = args.auto

driver = get_chrome_webdriver()
scraper = Scraper(driver, url, StreamServers.YOURUPLOAD)

# Parsing start and end points
start = args.start
end = args.end
start_given = False if type(start).__name__ == "NoneType" else True
end_given = False if type(end).__name__ == "NoneType" else True

if (start_given and start < 0) or (end_given and end < 0) or ((start_given and end_given) and (end < start)):
    print("Error: Invalid start and end points")
    sys.exit()

if start_given and end_given:
    episodes = list(range(start, end+1))
elif start_given and not end_given:
    # get episode numbers
    # make list from given start point to end point
    pass
elif not start_given and end_given:
    episodes = list(range(1,end+1))

# To find missing

# Check the available episodes in directory (or in the JSON file depending on options)
# Make a list of the missing episodes and fetch

# Parse user input
# Options: Anime URL, Start episode, End episode, Missing -> [Use metadata, Use file list], Auto-add, Metadata only
# Read config
# Check URL and get matching scraper instance
# Fetch episodes from prefered server
# Prepare JSON and simple text files

fetched_episodes = {}

# meta_data = scraper.fetch_metadata() # -> { title: _____, synopsis: _____, poster: ____ }

if episodes:
    for ep_nuber in episodes:
        fetched_episodes["Episode " + str(ep_nuber)] = scraper.fetch_episode_number(ep_nuber)
else:
    scraper.fetch_all_episodes(fetched_episodes)

driver.close()

print(fetched_episodes)

# Output operable data
<|endoftext|>


<commit_before>import sys
import time

import geopy.exc
from geopy.geocoders import GoogleV3
from joblib import Memory
import requests

from dataset import Database, Events
import settings

memory = Memory('cache/scrape', verbose=0)

def rate_limit_sleep(seconds):
    minutes = round(seconds / 60, 1)
    print('!', 'Waiting', minutes, 'minutes.')
    time.sleep(seconds)

class GitHub:
    def __init__(self):
        self.client_id = settings.CLIENT_ID
        self.client_secret = settings.CLIENT_SECRET

    def get(self, url, params=None):
        if params is None:
            params = {}

        params['client_id'] = self.client_id
        params['client_secret'] = self.client_secret

        response = requests.get(url, params=params)

        headers = response.headers
        if headers['X-RateLimit-Remaining'] == '0':
            reset_time = int(headers['X-RateLimit-Reset'])
            seconds_remaining = reset_time - int(time.time())
            rate_limit_sleep(seconds_remaining)
            return self.get(url, params)

        return response

    def get_all_users(self, since=0):
        next_url = 'https://api.github.com/users?since={}'.format(since)

        while next_url is not None:
            response = self.get(next_url)
            next_url = response.links['next']['url']
            for user in response.json():
                yield user

    def get_user(self, username):
        url = 'https://api.github.com/users/{}'.format(username)
        return self.get(url).json()

    def get_following_users(self, username):
        next_url = 'https://api.github.com/users/{}/following'.format(username)

        while next_url is not None:
            response = self.get(next_url)

            try:
                next_url = response.links['next']['url']
            except KeyError:
                next_url = None

            for user in response.json():
                yield user

class Geography:
    def __init__(self):
        self.api_key = settings.GOOGLE_API_KEY
        self.geolocator = GoogleV3(self.api_key)

        self.geocode = memory.cache(self.geocode)

    def geocode(self, text):
        try:
            result = self.geolocator.geocode(text)
        except geopy.exc.GeocoderQuotaExceeded:
            rate_limit_sleep(60 * 60)
            return self.geocode(text)
        except geopy.exc.GeocoderTimedOut:
            rate_limit_sleep(10)
            return self.geocode(text)
        else:
            time.sleep(0.1)  # to avoid rate limiting
            return result

    @staticmethod
    def get_country(location):
        for component in location.raw['address_components']:
            if 'country' in component['types']:
                return component['short_name'], component['long_name']

        raise ValueError()

class Genderize:
    def __init__(self):
        self.api_key = settings.GENDERIZE_API_KEY
        self.guess = memory.cache(self.guess)

    def guess(self, name):
        if name == "<script>alert('test')</script>":
            return '?', None

        params = {'name': name, 'apikey': self.api_key}
        response = requests.get('http://api.genderize.io', params=params)

        headers = response.headers
        if 'X-Rate-Limit-Remaining' not in headers:
            rate_limit_sleep(60 * 60)
            return self.guess(name)

        if headers['X-Rate-Limit-Remaining'] == '0':
            reset_time = int(headers['X-Rate-Limit-Reset'])
            seconds_remaining = reset_time
            rate_limit_sleep(seconds_remaining)
            return self.guess(name)

        data = response.json()

        if data['gender'] is None:
            return '?', None

        probability = float(data['probability'])
        code = data['gender'][0].upper()
        return code, probability

def print_status(user_id, *args):
    print('>', '#{}'.format(user_id), *args)
    sys.stdout.flush()

def users():
    github = GitHub()
    database = Database()
    events = Events()

    for event in events.iterate():
        try:
            actor = event['actor']
            user_id = actor['id']
            user_login = actor['login']
        except KeyError:
            continue

        if database.has_user(user_id):
            continue

        github_user = github.get_user(user_login)
        if 'id' not in github_user or user_id != github_user['id']:  # no longer a user
            database.insert_user({'id': user_id, 'deleted': True})
            print_status(user_id, ':(')
            continue

        fields = {
            'id': github_user['id'],
            'hireable': github_user['hireable']
        }

        for field in ['login', 'avatar_url', 'gravatar_id', 'name', 'company',
                      'blog', 'location', 'email', 'bio']:
            if github_user[field] is None:
                fields[field] = None
            else:
                fields[field] = github_user[field].strip()

        database.insert_user(fields)

        print_status(fields['id'], fields['login'], github_user['created_at'], '✓')

    database.close()

def user_followings():
    github = GitHub()

    """for user in session.query(User):
        for other_user in github.get_following_users(user.login):
            print(other_user)"""

def user_locations():
    geography = Geography()
    database = Database()

    for user_id, location_str in database.users_without_location:
        location = geography.geocode(location_str)
        if location is None:
            continue

        try:
            country_code, country_name = geography.get_country(location)
        except ValueError:
            continue

        print_status(user_id, location_str, '->', country_code,
                     '({})'.format(country_name))

        database.update_user_location(user_id, location.latitude,
                                      location.longitude, country_code)

def user_genders():
    genderize = Genderize()
    database = Database()

    for user_id, name in database.users_without_gender:
        gender, probability = genderize.guess(name.split()[0])

        try:
            probability_str = '({}%)'.format(probability * 100)
        except TypeError:
            probability_str = '(N/A)'

        print_status(user_id, name, '->', gender, probability_str)

        database.update_user_gender(user_id, gender, probability)

if __name__ == '__main__':
    for name in sys.argv[1:]:
        locals()[name]()
<commit_msg>Include login in non-existent users<commit_after>import sys
import time

import geopy.exc
from geopy.geocoders import GoogleV3
from joblib import Memory
import requests

from dataset import Database, Events
import settings

memory = Memory('cache/scrape', verbose=0)

def rate_limit_sleep(seconds):
    minutes = round(seconds / 60, 1)
    print('!', 'Waiting', minutes, 'minutes.')
    time.sleep(seconds)

class GitHub:
    def __init__(self):
        self.client_id = settings.CLIENT_ID
        self.client_secret = settings.CLIENT_SECRET

    def get(self, url, params=None):
        if params is None:
            params = {}

        params['client_id'] = self.client_id
        params['client_secret'] = self.client_secret

        response = requests.get(url, params=params)

        headers = response.headers
        if headers['X-RateLimit-Remaining'] == '0':
            reset_time = int(headers['X-RateLimit-Reset'])
            seconds_remaining = reset_time - int(time.time())
            rate_limit_sleep(seconds_remaining)
            return self.get(url, params)

        return response

    def get_all_users(self, since=0):
        next_url = 'https://api.github.com/users?since={}'.format(since)

        while next_url is not None:
            response = self.get(next_url)
            next_url = response.links['next']['url']
            for user in response.json():
                yield user

    def get_user(self, username):
        url = 'https://api.github.com/users/{}'.format(username)
        return self.get(url).json()

    def get_following_users(self, username):
        next_url = 'https://api.github.com/users/{}/following'.format(username)

        while next_url is not None:
            response = self.get(next_url)

            try:
                next_url = response.links['next']['url']
            except KeyError:
                next_url = None

            for user in response.json():
                yield user

class Geography:
    def __init__(self):
        self.api_key = settings.GOOGLE_API_KEY
        self.geolocator = GoogleV3(self.api_key)

        self.geocode = memory.cache(self.geocode)

    def geocode(self, text):
        try:
            result = self.geolocator.geocode(text)
        except geopy.exc.GeocoderQuotaExceeded:
            rate_limit_sleep(60 * 60)
            return self.geocode(text)
        except geopy.exc.GeocoderTimedOut:
            rate_limit_sleep(10)
            return self.geocode(text)
        else:
            time.sleep(0.1)  # to avoid rate limiting
            return result

    @staticmethod
    def get_country(location):
        for component in location.raw['address_components']:
            if 'country' in component['types']:
                return component['short_name'], component['long_name']

        raise ValueError()

class Genderize:
    def __init__(self):
        self.api_key = settings.GENDERIZE_API_KEY
        self.guess = memory.cache(self.guess)

    def guess(self, name):
        if name == "<script>alert('test')</script>":
            return '?', None

        params = {'name': name, 'apikey': self.api_key}
        response = requests.get('http://api.genderize.io', params=params)

        headers = response.headers
        if 'X-Rate-Limit-Remaining' not in headers:
            rate_limit_sleep(60 * 60)
            return self.guess(name)

        if headers['X-Rate-Limit-Remaining'] == '0':
            reset_time = int(headers['X-Rate-Limit-Reset'])
            seconds_remaining = reset_time
            rate_limit_sleep(seconds_remaining)
            return self.guess(name)

        data = response.json()

        if data['gender'] is None:
            return '?', None

        probability = float(data['probability'])
        code = data['gender'][0].upper()
        return code, probability

def print_status(user_id, *args):
    print('>', '#{}'.format(user_id), *args)
    sys.stdout.flush()

def users():
    github = GitHub()
    database = Database()
    events = Events()

    for event in events.iterate():
        try:
            actor = event['actor']
            user_id = actor['id']
            user_login = actor['login']
        except KeyError:
            continue

        if database.has_user(user_id):
            continue

        github_user = github.get_user(user_login)
        if 'id' not in github_user or user_id != github_user['id']:  # no longer a user
            database.insert_user({'id': user_id, 'login': user_login, 'deleted': True})
            print_status(user_id, ':(')
            continue

        fields = {
            'id': github_user['id'],
            'hireable': github_user['hireable']
        }

        for field in ['login', 'avatar_url', 'gravatar_id', 'name', 'company',
                      'blog', 'location', 'email', 'bio']:
            if github_user[field] is None:
                fields[field] = None
            else:
                fields[field] = github_user[field].strip()

        database.insert_user(fields)

        print_status(fields['id'], fields['login'], github_user['created_at'], '✓')

    database.close()

def user_followings():
    github = GitHub()

    """for user in session.query(User):
        for other_user in github.get_following_users(user.login):
            print(other_user)"""

def user_locations():
    geography = Geography()
    database = Database()

    for user_id, location_str in database.users_without_location:
        location = geography.geocode(location_str)
        if location is None:
            continue

        try:
            country_code, country_name = geography.get_country(location)
        except ValueError:
            continue

        print_status(user_id, location_str, '->', country_code,
                     '({})'.format(country_name))

        database.update_user_location(user_id, location.latitude,
                                      location.longitude, country_code)

def user_genders():
    genderize = Genderize()
    database = Database()

    for user_id, name in database.users_without_gender:
        gender, probability = genderize.guess(name.split()[0])

        try:
            probability_str = '({}%)'.format(probability * 100)
        except TypeError:
            probability_str = '(N/A)'

        print_status(user_id, name, '->', gender, probability_str)

        database.update_user_gender(user_id, gender, probability)

if __name__ == '__main__':
    for name in sys.argv[1:]:
        locals()[name]()
<|endoftext|>


<commit_before>"""
Build database model definitions
"""

# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.utils.translation import ugettext as _
from django.core.exceptions import ValidationError

from django.urls import reverse
from django.db import models
from django.core.validators import MinValueValidator

class Build(models.Model):
    """ A Build object organises the creation of new parts from the component parts.

    Attributes:
        part: The part to be built (from component BOM items)
        title: Brief title describing the build (required)
        quantity: Number of units to be built
        status: Build status code
        batch: Batch code transferred to build parts (optional)
        creation_date: Date the build was created (auto)
        completion_date: Date the build was completed
        URL: External URL for extra information
        notes: Text notes
    """

    def __str__(self):
        return "Build {q} x {part}".format(q=self.quantity, part=str(self.part))

    def get_absolute_url(self):
        return reverse('build-detail', kwargs={'pk': self.id})

    part = models.ForeignKey('part.Part', on_delete=models.CASCADE,
                             related_name='builds',
                             limit_choices_to={'buildable': True},
                             )

    title = models.CharField(max_length=100, help_text='Brief description of the build')

    quantity = models.PositiveIntegerField(
        default=1,
        validators=[MinValueValidator(1)],
        help_text='Number of parts to build'
    )

    # Build status codes
    PENDING = 10  # Build is pending / active
    HOLDING = 20  # Build is currently being held
    CANCELLED = 30  # Build was cancelled
    COMPLETE = 40  # Build is complete

    #: Build status codes
    BUILD_STATUS_CODES = {PENDING: _("Pending"),
                          HOLDING: _("Holding"),
                          CANCELLED: _("Cancelled"),
                          COMPLETE: _("Complete"),
                          }

    status = models.PositiveIntegerField(default=PENDING,
                                         choices=BUILD_STATUS_CODES.items(),
                                         validators=[MinValueValidator(0)])

    batch = models.CharField(max_length=100, blank=True, null=True,
                             help_text='Batch code for this build output')

    creation_date = models.DateField(auto_now=True, editable=False)

    completion_date = models.DateField(null=True, blank=True)

    URL = models.URLField(blank=True, help_text='Link to external URL')

    notes = models.TextField(blank=True)
    """ Notes attached to each build output """

    @property
    def required_parts(self):
        """ Returns a dict of parts required to build this part (BOM) """
        parts = []

        for item in self.part.bom_items.all():
            part = {'part': item.sub_part,
                    'per_build': item.quantity,
                    'quantity': item.quantity * self.quantity
                    }

            parts.append(part)

        return parts

    @property
    def can_build(self):
        """ Return true if there are enough parts to supply build """

        for item in self.required_parts:
            if item['part'].total_stock < item['quantity']:
                return False

        return True

    @property
    def is_active(self):
        """ Is this build active? An active build is either:

        - PENDING
        - HOLDING
        """

        return self.status in [
            self.PENDING,
            self.HOLDING
        ]

    @property
    def is_complete(self):
        """ Returns True if the build status is COMPLETE """
        return self.status == self.COMPLETE

class BuildItem(models.Model):
    """ A BuildItem links multiple StockItem objects to a Build.
    These are used to allocate part stock to a build.
    Once the Build is completed, the parts are removed from stock and the
    BuildItemAllocation objects are removed.

    Attributes:
        build: Link to a Build object
        stock: Link to a StockItem object
        quantity: Number of units allocated
    """

    def get_absolute_url(self):
        # TODO - Fix!
        return '/build/item/{pk}/'.format(pk=self.id)
        #return reverse('build-detail', kwargs={'pk': self.id})

    class Meta:
        unique_together = [
            ('build', 'stock_item'),
        ]

    def clean(self):
        """ Check validity of the BuildItem model.
        The following checks are performed:

        - StockItem.part must be in the BOM of the Part object referenced by Build
        """

        if self.stock_item.part not in self.build.part.required_parts():
            print('stock_item:', self.stock_item.part)
            for p in self.build.part.bom_items.all():
                print('bom_part:', p)
            raise ValidationError(
                {'stock_item': _("Selected stock item not found in BOM for part '{p}'".format(p=str(self.build.part)))}
            )

    build = models.ForeignKey(
        Build,
        on_delete=models.CASCADE,
        related_name='allocated_stock',
    )

    stock_item = models.ForeignKey(
        'stock.StockItem',
        on_delete=models.CASCADE,
        related_name='allocations',
        help_text='Stock Item to allocate to build',
    )

    quantity = models.PositiveIntegerField(
        default=1,
        validators=[MinValueValidator(1)],
        help_text='Stock quantity to allocate to build'
    )
<commit_msg>Catch BuildItem errors in clean()<commit_after>"""
Build database model definitions
"""

# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.utils.translation import ugettext as _
from django.core.exceptions import ValidationError

from django.urls import reverse
from django.db import models
from django.core.validators import MinValueValidator

class Build(models.Model):
    """ A Build object organises the creation of new parts from the component parts.

    Attributes:
        part: The part to be built (from component BOM items)
        title: Brief title describing the build (required)
        quantity: Number of units to be built
        status: Build status code
        batch: Batch code transferred to build parts (optional)
        creation_date: Date the build was created (auto)
        completion_date: Date the build was completed
        URL: External URL for extra information
        notes: Text notes
    """

    def __str__(self):
        return "Build {q} x {part}".format(q=self.quantity, part=str(self.part))

    def get_absolute_url(self):
        return reverse('build-detail', kwargs={'pk': self.id})

    part = models.ForeignKey('part.Part', on_delete=models.CASCADE,
                             related_name='builds',
                             limit_choices_to={'buildable': True},
                             )

    title = models.CharField(max_length=100, help_text='Brief description of the build')

    quantity = models.PositiveIntegerField(
        default=1,
        validators=[MinValueValidator(1)],
        help_text='Number of parts to build'
    )

    # Build status codes
    PENDING = 10  # Build is pending / active
    HOLDING = 20  # Build is currently being held
    CANCELLED = 30  # Build was cancelled
    COMPLETE = 40  # Build is complete

    #: Build status codes
    BUILD_STATUS_CODES = {PENDING: _("Pending"),
                          HOLDING: _("Holding"),
                          CANCELLED: _("Cancelled"),
                          COMPLETE: _("Complete"),
                          }

    status = models.PositiveIntegerField(default=PENDING,
                                         choices=BUILD_STATUS_CODES.items(),
                                         validators=[MinValueValidator(0)])

    batch = models.CharField(max_length=100, blank=True, null=True,
                             help_text='Batch code for this build output')

    creation_date = models.DateField(auto_now=True, editable=False)

    completion_date = models.DateField(null=True, blank=True)

    URL = models.URLField(blank=True, help_text='Link to external URL')

    notes = models.TextField(blank=True)
    """ Notes attached to each build output """

    @property
    def required_parts(self):
        """ Returns a dict of parts required to build this part (BOM) """
        parts = []

        for item in self.part.bom_items.all():
            part = {'part': item.sub_part,
                    'per_build': item.quantity,
                    'quantity': item.quantity * self.quantity
                    }

            parts.append(part)

        return parts

    @property
    def can_build(self):
        """ Return true if there are enough parts to supply build """

        for item in self.required_parts:
            if item['part'].total_stock < item['quantity']:
                return False

        return True

    @property
    def is_active(self):
        """ Is this build active? An active build is either:

        - PENDING
        - HOLDING
        """

        return self.status in [
            self.PENDING,
            self.HOLDING
        ]

    @property
    def is_complete(self):
        """ Returns True if the build status is COMPLETE """
        return self.status == self.COMPLETE

class BuildItem(models.Model):
    """ A BuildItem links multiple StockItem objects to a Build.
    These are used to allocate part stock to a build.
    Once the Build is completed, the parts are removed from stock and the
    BuildItemAllocation objects are removed.

    Attributes:
        build: Link to a Build object
        stock: Link to a StockItem object
        quantity: Number of units allocated
    """

    def get_absolute_url(self):
        # TODO - Fix!
        return '/build/item/{pk}/'.format(pk=self.id)
        #return reverse('build-detail', kwargs={'pk': self.id})

    class Meta:
        unique_together = [
            ('build', 'stock_item'),
        ]

    def clean(self):
        """ Check validity of the BuildItem model.
        The following checks are performed:

        - StockItem.part must be in the BOM of the Part object referenced by Build
        - Allocation quantity cannot exceed available quantity
        """

        super(BuildItem, self).clean()

        errors = {}

        if self.stock_item is not None and self.stock_item.part is not None:
            if self.stock_item.part not in self.build.part.required_parts():
                errors['stock_item'] = _("Selected stock item not found in BOM for part '{p}'".format(p=str(self.build.part)))

        if self.stock_item is not None and self.quantity > self.stock_item.quantity:
            errors['quantity'] = _("Allocated quantity ({n}) must not exceed available quantity ({q})".format(
                n=self.quantity,
                q=self.stock_item.quantity
            ))

        if len(errors) > 0:
            raise ValidationError(errors)

    build = models.ForeignKey(
        Build,
        on_delete=models.CASCADE,
        related_name='allocated_stock',
    )

    stock_item = models.ForeignKey(
        'stock.StockItem',
        on_delete=models.CASCADE,
        related_name='allocations',
        help_text='Stock Item to allocate to build',
    )

    quantity = models.PositiveIntegerField(
        default=1,
        validators=[MinValueValidator(1)],
        help_text='Stock quantity to allocate to build'
    )
<|endoftext|>


<commit_before>import discord
import asyncio
from tqdm import tqdm
import argparse

parser = argparse.ArgumentParser(description='Discord channel scraper')
requiredNamed = parser.add_argument_group('Required arguments:')
requiredNamed.add_argument('-c', '--channel', type=str, help='Channel to scrape. Requires the channel ID.', required=True)
requiredNamed.add_argument('-o', '--output', type=str, help='Output file in form *.txt. Will be stored in the same directory.', required=True)

args = parser.parse_args()

client = discord.Client()

@client.event
async def on_ready():
    print('Connection successful.')
    print('Your ID: ' + client.user.id)
    target = open(args.output, 'w')
    print(args.output, 'has been opened.')

    messageCount = 0
    channel = discord.Object(id=args.channel)

    print("Scraping messages... Don't send any messages while scraping!")
    with tqdm(leave=True,unit=' messages') as scraped:
        async for msg in client.logs_from(channel, 10000000000):
            line = "{} - {} - {}".format(msg.timestamp,msg.author.name, msg.content)
            line = line.encode('utf-8')
            toWrite = "{}".format(line)
            target.write(toWrite)
            target.write("\n")
            messageCount += 1
            scraped.update(1)
    print('-----')
    print('Scraping complete.')

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

client.run('email', 'password')
<commit_msg>Update to use token instead of email+pass<commit_after>import discord
import asyncio
from tqdm import tqdm
import argparse

parser = argparse.ArgumentParser(description='Discord channel scraper')
requiredNamed = parser.add_argument_group('Required arguments:')
requiredNamed.add_argument('-c', '--channel', type=str, help='Channel to scrape. Requires the channel ID.', required=True)
requiredNamed.add_argument('-o', '--output', type=str, help='Output file in form *.txt. Will be stored in the same directory.', required=True)

args = parser.parse_args()

client = discord.Client()

@client.event
async def on_ready():
    print('Connection successful.')
    print('Your ID: ' + client.user.id)
    target = open(args.output, 'w')
    print(args.output, 'has been opened.')

    messageCount = 0
    channel = discord.Object(id=args.channel)

    print("Scraping messages... Don't send any messages while scraping!")
    with tqdm(leave=True,unit=' messages') as scraped:
        async for msg in client.logs_from(channel, 10000000000):
            line = "{} - {} - {}".format(msg.timestamp,msg.author.name, msg.content)
            line = line.encode('utf-8')
            toWrite = "{}".format(line)
            target.write(toWrite)
            target.write("\n")
            messageCount += 1
            scraped.update(1)
    print('-----')
    print('Scraping complete.')

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

client.run(usertoken,bot=False)
<|endoftext|>


<commit_before># Copyright 2014 tsuru-circus authors. All rights reserved.
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file.

import json
import re
import requests
import logging

from logging.handlers import SysLogHandler
from socket import SOCK_DGRAM, SOCK_STREAM, gethostname
from tsuru import common
from datetime import datetime

def extract_message(msg):
    # 2012-11-06 18:30:10 [13887] [INFO]
    regex = "\d+\-\d+\-\d+ \d+\:\d+\:\d+ \[\d+\] \[\w+\] "
    msgs = re.split(regex, msg)
    return [m for m in msgs if m]

class Stream(object):

    def __init__(self, **kwargs):
        self._buffer = ""
        self._max_buffer_size = kwargs.get("max_buffer_size", 10240)
        self.apprc = "/home/application/apprc"
        self.watcher_name = kwargs.get("watcher_name", "")
        self.timeout = kwargs.get("timeout", 2)
        self.hostname = gethostname()

    def __call__(self, data):
        (appname, host, token, syslog_server, syslog_port,
         syslog_facility, syslog_socket) = self.load_envs()
        if appname and host and token:
            self.log_tsuru_api(data, appname, host, token)
        if syslog_server and syslog_port and syslog_facility:
            self.log_syslog(data, appname, syslog_server, syslog_port,
                            syslog_facility, syslog_socket)

    def log_tsuru_api(self, data, appname, host, token):
        url = "{0}/apps/{1}/log?source={2}&unit={3}".format(host, appname,
                                                            self.watcher_name,
                                                            self.hostname)
        messages = self.get_messages(data["data"])
        try:
            requests.post(url, data=json.dumps(messages),
                          headers={"Authorization": "bearer " + token},
                          timeout=self.timeout)
        except:
            pass

    def log_syslog(self, data, appname, host, port, facility, socket):
        messages = self.get_messages(data["data"])
        if socket == 'tcp':
            socket_type = SOCK_STREAM
        else:
            socket_type = SOCK_DGRAM
        try:
            date_utc = datetime.utcnow().isoformat("T") + "Z"
            logger = logging.getLogger(appname)
            logger.handlers = []
            logger.setLevel(logging.INFO)
            syslog = SysLogHandler(address=(host, int(port)),
                                   facility=facility, socktype=socket_type)
            formatter = logging.Formatter('{0} {1} %(name)s:\
                                           %(message)s'.format(date_utc,
                                                               self.hostname)),
            syslog.setFormatter(formatter)
            logger.addHandler(syslog)
            for m in messages:
                if data["name"] == 'stdout':
                    logger.info(m)
                else:
                    logger.error(m)
        except:
            pass

    def load_envs(self):
        envs = common.load_envs(self.apprc)
        return (envs.get("TSURU_APPNAME"), envs.get("TSURU_HOST"),
                envs.get("TSURU_APP_TOKEN"), envs.get("TSURU_SYSLOG_SERVER"),
                envs.get("TSURU_SYSLOG_PORT"),
                envs.get("TSURU_SYSLOG_FACILITY"),
                envs.get("TSURU_SYSLOG_SOCKET"))

    def close(self):
        pass

    def get_messages(self, msg):
        result = []
        if self._buffer != "":
            msg = self._buffer + msg
            self._buffer = ""
        msgs = extract_message(msg)
        lines = "".join(msgs).splitlines(True)
        for line in lines:
            if (line.endswith("\n") or len(line) > self._max_buffer_size):
                result.append(line)
            else:
                self._buffer = line
        return result
<commit_msg>Revert "fix syslog datetime field to UTC"<commit_after># Copyright 2014 tsuru-circus authors. All rights reserved.
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file.

import json
import re
import requests
import logging

from logging.handlers import SysLogHandler
from socket import SOCK_DGRAM, SOCK_STREAM, gethostname
from tsuru import common

def extract_message(msg):
    # 2012-11-06 18:30:10 [13887] [INFO]
    regex = "\d+\-\d+\-\d+ \d+\:\d+\:\d+ \[\d+\] \[\w+\] "
    msgs = re.split(regex, msg)
    return [m for m in msgs if m]

class Stream(object):

    def __init__(self, **kwargs):
        self._buffer = ""
        self._max_buffer_size = kwargs.get("max_buffer_size", 10240)
        self.apprc = "/home/application/apprc"
        self.watcher_name = kwargs.get("watcher_name", "")
        self.timeout = kwargs.get("timeout", 2)
        self.hostname = gethostname()

    def __call__(self, data):
        (appname, host, token, syslog_server, syslog_port,
         syslog_facility, syslog_socket) = self.load_envs()
        if appname and host and token:
            self.log_tsuru_api(data, appname, host, token)
        if syslog_server and syslog_port and syslog_facility:
            self.log_syslog(data, appname, syslog_server, syslog_port,
                            syslog_facility, syslog_socket)

    def log_tsuru_api(self, data, appname, host, token):
        url = "{0}/apps/{1}/log?source={2}&unit={3}".format(host, appname,
                                                            self.watcher_name,
                                                            self.hostname)
        messages = self.get_messages(data["data"])
        try:
            requests.post(url, data=json.dumps(messages),
                          headers={"Authorization": "bearer " + token},
                          timeout=self.timeout)
        except:
            pass

    def log_syslog(self, data, appname, host, port, facility, socket):
        messages = self.get_messages(data["data"])
        if socket == 'tcp':
            socket_type = SOCK_STREAM
        else:
            socket_type = SOCK_DGRAM
        try:
            logger = logging.getLogger(appname)
            logger.handlers = []
            logger.setLevel(logging.INFO)
            syslog = SysLogHandler(address=(host, int(port)),
                                   facility=facility, socktype=socket_type)
            formatter = logging.Formatter('%(asctime)s {} %(name)s:\
                                           %(message)s'.format(self.hostname),
                                          '%b  %d %H:%M:%S')
            syslog.setFormatter(formatter)
            logger.addHandler(syslog)
            for m in messages:
                if data["name"] == 'stdout':
                    logger.info(m)
                else:
                    logger.error(m)
        except:
            pass

    def load_envs(self):
        envs = common.load_envs(self.apprc)
        return (envs.get("TSURU_APPNAME"), envs.get("TSURU_HOST"),
                envs.get("TSURU_APP_TOKEN"), envs.get("TSURU_SYSLOG_SERVER"),
                envs.get("TSURU_SYSLOG_PORT"),
                envs.get("TSURU_SYSLOG_FACILITY"),
                envs.get("TSURU_SYSLOG_SOCKET"))

    def close(self):
        pass

    def get_messages(self, msg):
        result = []
        if self._buffer != "":
            msg = self._buffer + msg
            self._buffer = ""
        msgs = extract_message(msg)
        lines = "".join(msgs).splitlines(True)
        for line in lines:
            if (line.endswith("\n") or len(line) > self._max_buffer_size):
                result.append(line)
            else:
                self._buffer = line
        return result
<|endoftext|>


<commit_before>###########################################################
# Demo WPS service for testing and debugging.
#
# See the werkzeug documentation on how to use the debugger:
# http://werkzeug.pocoo.org/docs/0.12/debug/
###########################################################

import os
import psutil
import click
from jinja2 import Environment, PackageLoader, select_autoescape
from pywps import configuration

from . import wsgi
from six.moves.urllib.parse import urlparse

PID_FILE = os.path.abspath(os.path.join(os.path.curdir, "pywps.pid"))

CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])

template_env = Environment(
    loader=PackageLoader('emu', 'templates'),
    autoescape=select_autoescape(['yml', 'xml'])
)

def write_user_config(**kwargs):
    config_templ = template_env.get_template('pywps.cfg')
    rendered_config = config_templ.render(**kwargs)
    config_file = os.path.abspath(os.path.join(os.path.curdir, ".custom.cfg"))
    with open(config_file, 'w') as fp:
        fp.write(rendered_config)
    return config_file

def get_host():
    url = configuration.get_config_value('server', 'url')
    url = url or 'http://localhost:5000/wps'

    click.echo("starting WPS service on {}".format(url))

    parsed_url = urlparse(url)
    if ':' in parsed_url.netloc:
        host, port = parsed_url.netloc.split(':')
        port = int(port)
    else:
        host = parsed_url.netloc
        port = 80
    return host, port

def run_process_action(action=None):
    """Run an action with psutil on current process
    and return a status message."""
    action = action or 'status'
    try:
        with open(PID_FILE, 'r') as fp:
            pid = int(fp.read())
            p = psutil.Process(pid)
            if action == 'stop':
                p.terminate()
                msg = "pid={}, status=terminated".format(p.pid)
            else:
                from psutil import _pprint_secs
                msg = "pid={}, status={}, created={}".format(
                    p.pid, p.status(), _pprint_secs(p.create_time()))
        if action == 'stop':
            os.remove(PID_FILE)
    except IOError:
        msg = "No PID file found. Service not running?"
    except psutil.NoSuchProcess as e:
        msg = e.msg
    click.echo(msg)

def _run(application, bind_host=None, daemon=False):
    from werkzeug.serving import run_simple
    # call this *after* app is initialized ... needs pywps config.
    host, port = get_host()
    bind_host = bind_host or host
    # need to serve the wps outputs
    static_files = {
        '/outputs': configuration.get_config_value('server', 'outputpath')
    }
    run_simple(
        hostname=bind_host,
        port=port,
        application=application,
        use_debugger=False,
        use_reloader=False,
        threaded=True,
        # processes=2,
        use_evalex=not daemon,
        static_files=static_files)

@click.group(context_settings=CONTEXT_SETTINGS)
@click.version_option()
def cli():
    """Command line to start/stop a PyWPS service.

    Do not use this service in a production environment.
    It's intended to be running in a test environment only!
    For more documentation, visit http://pywps.org/doc
    """
    pass

@cli.command()
def status():
    """Show status of PyWPS service"""
    run_process_action(action='status')

@cli.command()
def stop():
    """Stop PyWPS service"""
    run_process_action(action='stop')

@cli.command()
@click.option('--config', '-c', metavar='PATH', help='path to pywps configuration file.')
@click.option('--bind-host', '-b', metavar='IP-ADDRESS', default='0.0.0.0',
              help='IP address used to bind service.')
@click.option('--daemon', '-d', is_flag=True, help='run in daemon mode.')
@click.option('--hostname', metavar='HOSTNAME', default='localhost', help='hostname in PyWPS configuration.')
@click.option('--port', metavar='PORT', default='5000', help='port in PyWPS configuration.')
@click.option('--maxsingleinputsize', default='200mb', help='maxsingleinputsize in PyWPS configuration.')
@click.option('--maxprocesses', metavar='INT', default='10', help='maxprocesses in PyWPS configuration.')
@click.option('--parallelprocesses', metavar='INT', default='2', help='parallelprocesses in PyWPS configuration.')
@click.option('--log-level', metavar='LEVEL', default='INFO', help='log level in PyWPS configuration.')
@click.option('--log-file', metavar='PATH', default='pywps.log', help='log file in PyWPS configuration.')
@click.option('--database', default='sqlite:///pywps-logs.sqlite', help='database in PyWPS configuration')
def start(config, bind_host, daemon, hostname, port,
          maxsingleinputsize, maxprocesses, parallelprocesses,
          log_level, log_file, database):
    """Start PyWPS service.
    This service is by default available at http://localhost:5000/wps
    """
    cfgfiles = []
    cfgfiles.append(write_user_config(
        wps_hostname=hostname,
        wps_port=port,
        wps_maxsingleinputsize=maxsingleinputsize,
        wps_maxprocesses=maxprocesses,
        wps_parallelprocesses=parallelprocesses,
        wps_log_level=log_level,
        wps_log_file=log_file,
        wps_database=database,
    ))
    if config:
        cfgfiles.append(config)
    app = wsgi.create_app(cfgfiles)
    # let's start the service ...
    # See:
    # * https://github.com/geopython/pywps-flask/blob/master/demo.py
    # * http://werkzeug.pocoo.org/docs/0.14/serving/
    if daemon:
        # daemon (fork) mode
        pid = None
        try:
            pid = os.fork()
            if pid:
                click.echo('forked process id: {}'.format(pid))
                with open(PID_FILE, 'w') as fp:
                    fp.write("{}".format(pid))
        except OSError as e:
            raise Exception("%s [%d]" % (e.strerror, e.errno))

        if pid == 0:
            os.setsid()
            _run(app, bind_host=bind_host, daemon=True)
        else:
            os._exit(0)
    else:
        # no daemon
        _run(app, bind_host=bind_host)
<commit_msg>Fix codacy issues: autoescape=True<commit_after>###########################################################
# Demo WPS service for testing and debugging.
#
# See the werkzeug documentation on how to use the debugger:
# http://werkzeug.pocoo.org/docs/0.12/debug/
###########################################################

import os
import psutil
import click
from jinja2 import Environment, PackageLoader, select_autoescape
from pywps import configuration

from . import wsgi
from six.moves.urllib.parse import urlparse

PID_FILE = os.path.abspath(os.path.join(os.path.curdir, "pywps.pid"))

CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])

template_env = Environment(
    loader=PackageLoader('emu', 'templates'),
    autoescape=True
)

def write_user_config(**kwargs):
    config_templ = template_env.get_template('pywps.cfg')
    rendered_config = config_templ.render(**kwargs)
    config_file = os.path.abspath(os.path.join(os.path.curdir, ".custom.cfg"))
    with open(config_file, 'w') as fp:
        fp.write(rendered_config)
    return config_file

def get_host():
    url = configuration.get_config_value('server', 'url')
    url = url or 'http://localhost:5000/wps'

    click.echo("starting WPS service on {}".format(url))

    parsed_url = urlparse(url)
    if ':' in parsed_url.netloc:
        host, port = parsed_url.netloc.split(':')
        port = int(port)
    else:
        host = parsed_url.netloc
        port = 80
    return host, port

def run_process_action(action=None):
    """Run an action with psutil on current process
    and return a status message."""
    action = action or 'status'
    try:
        with open(PID_FILE, 'r') as fp:
            pid = int(fp.read())
            p = psutil.Process(pid)
            if action == 'stop':
                p.terminate()
                msg = "pid={}, status=terminated".format(p.pid)
            else:
                from psutil import _pprint_secs
                msg = "pid={}, status={}, created={}".format(
                    p.pid, p.status(), _pprint_secs(p.create_time()))
        if action == 'stop':
            os.remove(PID_FILE)
    except IOError:
        msg = "No PID file found. Service not running?"
    except psutil.NoSuchProcess as e:
        msg = e.msg
    click.echo(msg)

def _run(application, bind_host=None, daemon=False):
    from werkzeug.serving import run_simple
    # call this *after* app is initialized ... needs pywps config.
    host, port = get_host()
    bind_host = bind_host or host
    # need to serve the wps outputs
    static_files = {
        '/outputs': configuration.get_config_value('server', 'outputpath')
    }
    run_simple(
        hostname=bind_host,
        port=port,
        application=application,
        use_debugger=False,
        use_reloader=False,
        threaded=True,
        # processes=2,
        use_evalex=not daemon,
        static_files=static_files)

@click.group(context_settings=CONTEXT_SETTINGS)
@click.version_option()
def cli():
    """Command line to start/stop a PyWPS service.

    Do not use this service in a production environment.
    It's intended to be running in a test environment only!
    For more documentation, visit http://pywps.org/doc
    """
    pass

@cli.command()
def status():
    """Show status of PyWPS service"""
    run_process_action(action='status')

@cli.command()
def stop():
    """Stop PyWPS service"""
    run_process_action(action='stop')

@cli.command()
@click.option('--config', '-c', metavar='PATH', help='path to pywps configuration file.')
@click.option('--bind-host', '-b', metavar='IP-ADDRESS', default='127.0.0.1',
              help='IP address used to bind service.')
@click.option('--daemon', '-d', is_flag=True, help='run in daemon mode.')
@click.option('--hostname', metavar='HOSTNAME', default='localhost', help='hostname in PyWPS configuration.')
@click.option('--port', metavar='PORT', default='5000', help='port in PyWPS configuration.')
@click.option('--maxsingleinputsize', default='200mb', help='maxsingleinputsize in PyWPS configuration.')
@click.option('--maxprocesses', metavar='INT', default='10', help='maxprocesses in PyWPS configuration.')
@click.option('--parallelprocesses', metavar='INT', default='2', help='parallelprocesses in PyWPS configuration.')
@click.option('--log-level', metavar='LEVEL', default='INFO', help='log level in PyWPS configuration.')
@click.option('--log-file', metavar='PATH', default='pywps.log', help='log file in PyWPS configuration.')
@click.option('--database', default='sqlite:///pywps-logs.sqlite', help='database in PyWPS configuration')
def start(config, bind_host, daemon, hostname, port,
          maxsingleinputsize, maxprocesses, parallelprocesses,
          log_level, log_file, database):
    """Start PyWPS service.
    This service is by default available at http://localhost:5000/wps
    """
    cfgfiles = []
    cfgfiles.append(write_user_config(
        wps_hostname=hostname,
        wps_port=port,
        wps_maxsingleinputsize=maxsingleinputsize,
        wps_maxprocesses=maxprocesses,
        wps_parallelprocesses=parallelprocesses,
        wps_log_level=log_level,
        wps_log_file=log_file,
        wps_database=database,
    ))
    if config:
        cfgfiles.append(config)
    app = wsgi.create_app(cfgfiles)
    # let's start the service ...
    # See:
    # * https://github.com/geopython/pywps-flask/blob/master/demo.py
    # * http://werkzeug.pocoo.org/docs/0.14/serving/
    if daemon:
        # daemon (fork) mode
        pid = None
        try:
            pid = os.fork()
            if pid:
                click.echo('forked process id: {}'.format(pid))
                with open(PID_FILE, 'w') as fp:
                    fp.write("{}".format(pid))
        except OSError as e:
            raise Exception("%s [%d]" % (e.strerror, e.errno))

        if pid == 0:
            os.setsid()
            _run(app, bind_host=bind_host, daemon=True)
        else:
            os._exit(0)
    else:
        # no daemon
        _run(app, bind_host=bind_host)
<|endoftext|>


<commit_before># Copyright 2016 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.
"""Task manager for Turbinia."""

import logging
import time

import psq
from google.cloud import datastore
from google.cloud import pubsub

import turbinia
from turbinia import evidence
from turbinia import config
from turbinia import jobs
from turbinia import pubsub as turbinia_pubsub

def get_task_manager():
  config.LoadConfig()
  if config.TASK_MANAGER == 'PSQ':
    return PSQTaskManager()
  else:
    msg = u'Task Manager type "{0:s}" not implemented'.format(
        config.TASK_MANAGER)
    raise turbinia.TurbiniaException(msg)

def task_runner(obj, *args, **kwargs):
  """Wrapper function to run specified TurbiniaTask object.

  Args:
    obj: An instantiated TurbiniaTask object.
    *args: Any Args to pass to obj.
    **kwargs: Any keyword args to pass to obj.

  Returns:
    Output from TurbiniaTask (should be TurbiniaTaskReslt).
  """
  # TODO(aarontp): Add proper error checks/handling
  return obj.run(*args, **kwargs)

class TaskManager(object):
  """Class to manage Turbinia Tasks.

  Handles incoming new Evidence messages and adds new Tasks to the queue.

  Attributes:
    jobs: A list of registered and instantiated job objects
    evidence: A list of evidence objects to process
  """

  def __init__(self):
    self.jobs = []
    self.evidence = []

  def setup(self):
    """Does setup of Task manager and its dependencies."""
    self._backend_setup()
    self.jobs = jobs.GetJobs()

  def add_evidence(self, evidence_):
    """Adds new evidence and creates tasks to process it.

    This creates all tasks configured to process the given type of evidence.

    Args:
      evidence_: evidence object to add.
    """
    if not self.jobs:
      raise turbinia.TurbiniaException(
          u'Jobs must be registered before evidence can be added')
    logging.info(u'Adding new evidence: {0:s}'.format(str(evidence_)))
    self.evidence.append(evidence_)
    for job in self.jobs:
      if [True for t in job.evidence_input if isinstance(evidence_, t)]:
        logging.info(u'Adding {0:s} job to process {1:s}'.format(
            job.name, evidence_.name))
        # pylint: disable=expression-not-assigned
        [self.add_task(t) for t in job.create_tasks([evidence_])]

  def get_evidence(self):
    """Checks for new evidence.

    Returns:
      A list of Evidence objects
    """
    raise NotImplementedError

  def add_job(self, job):
    # TODO(aarontp): Insert jobs according to priority
    self.jobs.append(job)

  def get_job(self, job_id):
    """Returns an active job instance for a given job id.

    Args:
      job_id: The id of the job to search for.

    Returns:
      Job object if job is found, else None.
    """
    for job in self.jobs:
      if self.job.id == job_id:
        return job
    return None

  def add_task(self, task=None, evidence_=None):
    """Adds a task to be queued along with the evidence it will process.

    Args:
      task: An instantiated Turbinia Task
      evidence_: An Evidence object to be processed.
    """
    raise NotImplementedError

  def process_tasks(self):
    """Process any tasks that need to be processed.

    Returns:
      The number of tasks that have completed.
    """
    raise NotImplementedError

  def run(self):
    while True:
      # pylint: disable=expression-not-assigned
      [self.add_evidence(x) for x in self.get_evidence()]
      self.process_tasks()
      # TODO(aarontp): Add config var for this.
      time.sleep(30)

class PSQTaskManager(TaskManager):
  """PSQ implementation of TaskManager.

  Attributes:
    psq: PSQ Queue object.
    server_pubsub: A PubSubClient object for receiving new evidence messages.
    task_results: A list of outstanding PSQ task results.
  """

  def __init__(self):
    self.task_results = []
    config.LoadConfig()
    super(PSQTaskManager, self).__init__()

  def _backend_setup(self):
    """Set up backend dependencies."""
    self.server_pubsub = turbinia_pubsub.PubSubClient(config.PUBSUB_TOPIC)
    psq_pubsub_client = pubsub.Client(project=config.PROJECT)
    datastore_client = datastore.Client(project=config.PROJECT)
    self.psq = psq.Queue(
        psq_pubsub_client, config.PSQ_TOPIC,
        storage=psq.DatastoreStorage(datastore_client))

  def _complete_task(self, psq_task, task):
    """Runs final task data recording.

    Args:
      psq_task: An instance of the psq_task that ran
      task: The Turbinia Task object
    """
    # TODO(aarontp): Make sure this is set by the task
    if not task.result.successful:
      logging.error('Task {0:s} was not succesful'.format(task.name))
    else:
      logging.info('Task {0:s} executed with status {1:d}'.format(
          task.name, task.result))

    # Add output as new evidence to process
    if not task.output:
      logging.info('Task {0:s} did not return output'.format(task.name))
    elif isinstance(task.output, evidence.Evidence):
      logging.info('Task {0:s} returned non-Evidence output type {1:s}'.format(
          task.name, type(task.output)))
    else:
      self.add_evidence(task.output)

  def process_tasks(self):
    completed_tasks = []
    for result in self.task_results:
      psq_task = result.get_task()
      if not psq_task:
        logging.debug('Task {0:d} not yet created'.format(result.task_id))
      elif psq_task.status not in (psq.task.FINISHED, psq.task.FAILED):
        logging.debug('Task {0:d} still running'.format(psq_task.id))
      elif psq_task.status == psq.task.FAILED:
        logging.debug('Task {0:d} failed.'.format(psq_task.id))
        # TODO(aarontp): handle failures
      else:
        output = result.result()
        completed_tasks.append(result)
        self._complete_task(psq_task, output)

    # pylint: disable=expression-not-assigned
    return len([self.task_results.pop(task) for task in completed_tasks])

  def get_evidence(self):
    # TODO(aarontp): code goes here.
    pass

  def add_task(self, task, evidence_):
    logging.info('Adding task {0:s} with evidence {1:s} to queue').format(
        task.name, evidence_.name)
    self.task_results.append(self.psq.enqueue(task_runner, task, evidence_))

class PubSubTaskManager(TaskManager):

  def _process_task_message(self, message):
    """Process messages relating to task acceptance/update/completion."""
    # TODO(aarontp): fix
    if message[u'message_type'] == pubsub.TASKUPDATE:
      self._complete_task(message[u'job_id'], message[u'task_id'])
    elif message[u'message_type'] == pubsub.TASKSTART:
      self._complete_task(message[u'job_id'], message[u'task_id'])

  def add_worker(self, worker):
    self.worker.append(worker)

  def get_num_active_workers(self):
    return sum([1 for worker in self.workers if worker.in_use])

  def get_free_worker_count(self):
    return len(self.workers) - self.get_num_active_workers()

  def get_status(self):
    report_data = []
    report_data.append(u'Jobs:')
    report_data.append(u'\tName:\tActive Task:')
    for job in self.jobs:
      report_data.append(
          u'\t{0:s}\t{1:s}'.format(job.name, job.active_task.name))

    report_data.append(u'Workers:')
    report_data.append(u'\tId:\tHostname:\tActive Job:')
    for worker in self.workers:
      job = self.get_job(worker.active_job)
      job_name = u'No Active Job' if not job else job.name
      report_data.append(u'\t{0:s}\t{1:s}\t{2:s}'.format(
          worker.id, worker.hostname, job_name))

    return '\n'.join(report_data)

<commit_msg>Lots of fixes for TaskManager. First (mostly) working test end-to-end!<commit_after># Copyright 2016 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.
"""Task manager for Turbinia."""

import logging
import time

import psq
from google.cloud import datastore
from google.cloud import pubsub
from google.gax.errors import GaxError

import turbinia
from turbinia import evidence
from turbinia import config
from turbinia import jobs
from turbinia import pubsub as turbinia_pubsub

def get_task_manager():
  """Return task manager object based on config.

  Returns
    Initialized TaskManager object.
  """
  config.LoadConfig()
  if config.TASK_MANAGER == 'PSQ':
    return PSQTaskManager()
  else:
    msg = u'Task Manager type "{0:s}" not implemented'.format(
        config.TASK_MANAGER)
    raise turbinia.TurbiniaException(msg)

def task_runner(obj, *args, **kwargs):
  """Wrapper function to run specified TurbiniaTask object.

  Args:
    obj: An instantiated TurbiniaTask object.
    *args: Any Args to pass to obj.
    **kwargs: Any keyword args to pass to obj.

  Returns:
    Output from TurbiniaTask (should be TurbiniaTaskReslt).
  """
  # TODO(aarontp): Add proper error checks/handling
  return obj.run(*args, **kwargs)

class TaskManager(object):
  """Class to manage Turbinia Tasks.

  Handles incoming new Evidence messages and adds new Tasks to the queue.

  Attributes:
    jobs: A list of instantiated job objects
    evidence: A list of evidence objects to process
  """

  def __init__(self):
    self.jobs = []
    self.evidence = []

  def _backend_setup(self):
    """Sets up backend dependencies."""
    raise NotImplementedError

  def setup(self):
    """Does setup of Task manager and its dependencies."""
    self._backend_setup()
    # TODO(aarontp): Consider instantiating a job per evidence object
    self.jobs = jobs.get_jobs()

  def add_evidence(self, evidence_):
    """Adds new evidence and creates tasks to process it.

    This creates all tasks configured to process the given type of evidence.

    Args:
      evidence_: evidence object to add.
    """
    if not self.jobs:
      raise turbinia.TurbiniaException(
          u'Jobs must be registered before evidence can be added')
    logging.info(u'Adding new evidence: {0:s}'.format(str(evidence_)))
    self.evidence.append(evidence_)
    job_count = 0
    # TODO(aarontp): Add some kind of loop detection in here so that jobs can
    # register for Evidence(), or or other evidence types that may be a super
    # class of the output of the job itself.  Short term we could potentially
    # have a run time check for this upon Job instantiation to prevent it.
    for job in self.jobs:
      if [True for t in job.evidence_input if isinstance(evidence_, t)]:
        logging.info(u'Adding {0:s} job to process {1:s}'.format(
            job.name, evidence_.name))
        job_count += 1
        for task in job.create_tasks([evidence_]):
          self.add_task(task, evidence_)

    if not job_count:
      logging.warning('No Jobs/Tasks were created for Evidence [{0:s}]. '
                      'Jobs may need to be configured to allow this type of '
                      'Evidence as input'.format(evidence_.name))

  def get_evidence(self):
    """Checks for new evidence to process.

    Returns:
      A list of Evidence objects
    """
    raise NotImplementedError

  def add_task(self, task, evidence_):
    """Adds a task to be queued along with the evidence it will process.

    Args:
      task: An instantiated Turbinia Task
      evidence_: An Evidence object to be processed.
    """
    raise NotImplementedError

  def process_tasks(self):
    """Process any tasks that need to be processed.

    Returns:
      The number of tasks that have completed.
    """
    raise NotImplementedError

  def run(self):
    """Main run loop for TaskManager."""
    logging.info('Starting PSQ Task Manager run loop')
    # TODO(aarontp): Add early exit option.
    while True:
      # pylint: disable=expression-not-assigned
      [self.add_evidence(x) for x in self.get_evidence()]
      self.process_tasks()
      # TODO(aarontp): Add config var for this.
      time.sleep(30)

class PSQTaskManager(TaskManager):
  """PSQ implementation of TaskManager.

  Attributes:
    psq: PSQ Queue object.
    psq_task_results: A list of outstanding PSQ task results.
    server_pubsub: A PubSubClient object for receiving new evidence messages.
  """

  def __init__(self):
    self.psq = None
    self.psq_task_results = []
    self.server_pubsub = None
    config.LoadConfig()
    super(PSQTaskManager, self).__init__()

  def _backend_setup(self):
    self.server_pubsub = turbinia_pubsub.PubSubClient(config.PUBSUB_TOPIC)
    psq_pubsub_client = pubsub.Client(project=config.PROJECT)
    datastore_client = datastore.Client(project=config.PROJECT)
    try:
      self.psq = psq.Queue(
          psq_pubsub_client, config.PSQ_TOPIC,
          storage=psq.DatastoreStorage(datastore_client))
    except GaxError as e:
      msg = 'Error creating PSQ Queue: {0:s}'.format(str(e))
      logging.error(msg)
      raise turbinia.TurbiniaException(msg)

  def _finalize_result(self, task_result):
    """Runs final task results recording.

    self.process_tasks handles things that have failed at the PSQ layer, and
    this function handles tasks that have failed below that layer (i.e.
    somewhere in our Task code).

    Args:
      task_result: The TurbiniaTaskResult object
    """
    # TODO(aarontp): Make sure this is set by the task
    if not task_result.successful:
      logging.error(
          'Task {0:s} was not succesful'.format(task_result.task_name))
    else:
      logging.info('Task {0:s} executed with status {1:d}'.format(
          task_result.task_name, task_result.result))

    if not isinstance(task_result.evidence, list):
      logging.info(
          'Task {0:s} did not return list'.format(task_result.task_name))
      return

    for evidence_ in task_result.evidence:
      if isinstance(evidence_, evidence.Evidence):
        logging.info(u'Task {0:s} returned Evidence {1:s}'.format(
            task_result.task_name, evidence_.name))
        self.add_evidence(evidence_)
      else:
        logging.error(
            u'Task {0:s} returned non-Evidence output type {1:s}'.format(
                task_result.task_name, type(task_result.evidence)))

  def process_tasks(self):
    completed_tasks = []
    for psq_task_result in self.psq_task_results:
      psq_task = psq_task_result.get_task()
      # This handles tasks that have failed at the PSQ layer.
      if not psq_task:
        logging.debug(
            'Task {0:s} not yet created'.format(psq_task_result.task_id))
      elif psq_task.status not in (psq.task.FINISHED, psq.task.FAILED):
        logging.debug('Task {0:d} not finished'.format(psq_task.id))
      elif psq_task.status == psq.task.FAILED:
        logging.debug('Task {0:d} failed.'.format(psq_task.id))
        # TODO(aarontp): handle failures
      else:
        completed_tasks.append(psq_task_result)
        self._finalize_result(psq_task_result.result())

    # pylint: disable=expression-not-assigned
    return len([self.psq_task_results.remove(task) for task in completed_tasks])

  def get_evidence(self):
    # TODO(aarontp): code goes here.
    return []

  def add_task(self, task, evidence_):
    logging.info('Adding PSQ task {0:s} with evidence {1:s} to queue'.format(
        task.name, evidence_.name))
    self.psq_task_results.append(self.psq.enqueue(task_runner, task, evidence_))
<|endoftext|>


<commit_before># Copyright 2019 Tecnativa - Pedro M. Baeza
# License AGPL-3.0 or later (https://www.gnu.org/licenses/agpl).
{
    "name": "Employee Calendar Planning",
    "version": "14.0.1.0.1",
    "category": "Human Resources",
    "website": "https://github.com/OCA/hr",
    "author": "Tecnativa,Odoo Community Association (OCA)",
    "license": "AGPL-3",
    "installable": True,
    "depends": ["hr"],
    "data": ["security/ir.model.access.csv", "views/hr_employee_views.xml"],
    "post_init_hook": "post_init_hook",
}
<commit_msg>hr_employee_calendar_planning 14.0.1.0.2<commit_after># Copyright 2019 Tecnativa - Pedro M. Baeza
# License AGPL-3.0 or later (https://www.gnu.org/licenses/agpl).
{
    "name": "Employee Calendar Planning",
    "version": "14.0.1.0.2",
    "category": "Human Resources",
    "website": "https://github.com/OCA/hr",
    "author": "Tecnativa,Odoo Community Association (OCA)",
    "license": "AGPL-3",
    "installable": True,
    "depends": ["hr"],
    "data": ["security/ir.model.access.csv", "views/hr_employee_views.xml"],
    "post_init_hook": "post_init_hook",
}
<|endoftext|>