<commit_before>import pickle
import datetime
import pytest
import six
from pyrsistent import PRecord, field, InvariantException, ny, CheckedPSet, CheckedPVector

class ARecord(PRecord):
    x = field(type=(int, float))
    y = field()

def test_create():
    r = ARecord(x=1, y='foo')
    assert r.x == 1
    assert r.y == 'foo'
    assert isinstance(r, ARecord)

def test_correct_assignment():
    r = ARecord(x=1, y='foo')
    r2 = r.set('x', 2.0)
    r3 = r2.set('y', 'bar')

    assert r2 == {'x': 2.0, 'y': 'foo'}
    assert r3 == {'x': 2.0, 'y': 'bar'}
    assert isinstance(r3, ARecord)

def test_direct_assignment_not_possible():
    with pytest.raises(AttributeError):
        ARecord().x = 1

def test_cannot_assign_undeclared_fields():
    with pytest.raises(AttributeError):
        ARecord().set('z', 5)

def test_cannot_assign_wrong_type_to_fields():
    with pytest.raises(TypeError):
        ARecord().set('x', 'foo')

def test_cannot_construct_with_undeclared_fields():
    with pytest.raises(AttributeError):
        ARecord(z=5)

def test_cannot_construct_with_fields_of_wrong_type():
    with pytest.raises(TypeError):
        ARecord(x='foo')

def test_support_record_inheritance():
    class BRecord(ARecord):
        z = field()

    r = BRecord(x=1, y='foo', z='bar')

    assert isinstance(r, BRecord)
    assert isinstance(r, ARecord)
    assert r == {'x': 1, 'y': 'foo', 'z': 'bar'}

def test_single_type_spec():
    class A(PRecord):
        x = field(type=int)

    r = A(x=1)
    assert r.x == 1

    with pytest.raises(TypeError):
        r.set('x', 'foo')

def test_remove():
    r = ARecord(x=1, y='foo')
    r2 = r.remove('y')

    assert isinstance(r2, ARecord)
    assert r2 == {'x': 1}

def test_field_invariant_must_hold():
    class BRecord(PRecord):
        x = field(invariant=lambda x: (x > 1, 'x too small'))
        y = field(mandatory=True)

    try:
        BRecord(x=1)
        assert False
    except InvariantException as e:
        assert e.invariant_errors == ('x too small',)
        assert e.missing_fields == ('BRecord.y',)

def test_global_invariant_must_hold():
    class BRecord(PRecord):
        __invariant__ = lambda r: (r.x <= r.y, 'y smaller than x')
        x = field()
        y = field()

    BRecord(x=1, y=2)

    try:
        BRecord(x=2, y=1)
        assert False
    except InvariantException as e:
        assert e.invariant_errors == ('y smaller than x',)
        assert e.missing_fields == ()

def test_set_multiple_fields():
    a = ARecord(x=1, y='foo')
    b = a.set(x=2, y='bar')

    assert b == {'x': 2, 'y': 'bar'}

def test_initial_value():
    class BRecord(PRecord):
        x = field(initial=1)
        y = field(initial=2)

    a = BRecord()
    assert a.x == 1
    assert a.y == 2

def test_type_specification_must_be_a_type():
    with pytest.raises(TypeError):
        class BRecord(PRecord):
            x = field(type=1)

def test_initial_must_be_of_correct_type():
    with pytest.raises(TypeError):
        class BRecord(PRecord):
            x = field(type=int, initial='foo')

def test_invariant_must_be_callable():
    with pytest.raises(TypeError):
        class BRecord(PRecord):
            x = field(invariant='foo')

def test_global_invariants_are_inherited():
    class BRecord(PRecord):
        __invariant__ = lambda r: (r.x % r.y == 0, 'modulo')
        x = field()
        y = field()

    class CRecord(BRecord):
        __invariant__ = lambda r: (r.x > r.y, 'size')

    try:
        CRecord(x=5, y=3)
        assert False
    except InvariantException as e:
        assert e.invariant_errors == ('modulo',)

def test_global_invariants_must_be_callable():
    with pytest.raises(TypeError):
        class CRecord(PRecord):
            __invariant__ = 1

def test_repr():
    r = ARecord(x=1, y=2)
    assert repr(r) == 'ARecord(x=1, y=2)' or repr(r) == 'ARecord(y=2, x=1)'

def test_factory():
    class BRecord(PRecord):
        x = field(type=int, factory=int)

    assert BRecord(x=2.5) == {'x': 2}

def test_factory_must_be_callable():
    with pytest.raises(TypeError):
        class BRecord(PRecord):
            x = field(type=int, factory=1)

def test_nested_record_construction():
    class BRecord(PRecord):
        x = field(int, factory=int)

    class CRecord(PRecord):
        a = field()
        b = field(type=BRecord)

    r = CRecord.create({'a': 'foo', 'b': {'x': '5'}})
    assert isinstance(r, CRecord)
    assert isinstance(r.b, BRecord)
    assert r == {'a': 'foo', 'b': {'x': 5}}

def test_pickling():
    x = ARecord(x=2.0, y='bar')
    y = pickle.loads(pickle.dumps(x, -1))

    assert x == y
    assert isinstance(y, ARecord)

def test_all_invariant_errors_reported():
    class BRecord(PRecord):
        x = field(factory=int, invariant=lambda x: (x >= 0, 'x negative'))
        y = field(mandatory=True)

    class CRecord(PRecord):
        a = field(invariant=lambda x: (x != 0, 'a zero'))
        b = field(type=BRecord)

    try:
        CRecord.create({'a': 0, 'b': {'x': -5}})
        assert False
    except InvariantException as e:
        assert set(e.invariant_errors) == set(['x negative', 'a zero'])
        assert e.missing_fields == ('BRecord.y',)

def test_precord_factory_method_is_idempotent():
    class BRecord(PRecord):
        x = field()
        y = field()

    r = BRecord(x=1, y=2)
    assert BRecord.create(r) is r

def test_serialize():
    class BRecord(PRecord):
        d = field(type=datetime.date,
                  factory=lambda d: datetime.datetime.strptime(d, "%d%m%Y").date(),
                  serializer=lambda format, d: d.strftime('%Y-%m-%d') if format=='ISO' else d.strftime('%d%m%Y'))

    assert BRecord(d='14012015').serialize('ISO') == {'d': '2015-01-14'}
    assert BRecord(d='14012015').serialize('other') == {'d': '14012015'}

def test_nested_serialize():
    class BRecord(PRecord):
        d = field(serializer=lambda format, d: format)

    class CRecord(PRecord):
        b = field()

    serialized = CRecord(b=BRecord(d='foo')).serialize('bar')

    assert serialized == {'b': {'d': 'bar'}}
    assert isinstance(serialized, dict)

def test_serializer_must_be_callable():
    with pytest.raises(TypeError):
        class CRecord(PRecord):
            x = field(serializer=1)

def test_transform_without_update_returns_same_precord():
    r = ARecord(x=2.0, y='bar')
    assert r.transform([ny], lambda x: x) is r

class Application(PRecord):
    name = field(type=six.text_type)
    image = field(type=six.text_type)

class ApplicationVector(CheckedPVector):
    __type__ = Application

class Node(PRecord):
    applications = field(type=ApplicationVector)

def test_nested_create_serialize():
    node = Node(applications=[Application(name=u'myapp', image=u'myimage'),
                              Application(name=u'b', image=u'c')])

    node2 = Node.create({u'applications': [{u'name': u'myapp', u'image': u'myimage'},
                                           {u'name': u'b', u'image': u'c'}]})

    assert node == node2

    serialized = node.serialize()
    restored = Node.create(serialized)

    assert restored == node
<commit_msg>Python 3.2 compatibility for tests<commit_after>import pickle
import datetime
import pytest
import six
from pyrsistent import PRecord, field, InvariantException, ny, CheckedPSet, CheckedPVector

class ARecord(PRecord):
    x = field(type=(int, float))
    y = field()

def test_create():
    r = ARecord(x=1, y='foo')
    assert r.x == 1
    assert r.y == 'foo'
    assert isinstance(r, ARecord)

def test_correct_assignment():
    r = ARecord(x=1, y='foo')
    r2 = r.set('x', 2.0)
    r3 = r2.set('y', 'bar')

    assert r2 == {'x': 2.0, 'y': 'foo'}
    assert r3 == {'x': 2.0, 'y': 'bar'}
    assert isinstance(r3, ARecord)

def test_direct_assignment_not_possible():
    with pytest.raises(AttributeError):
        ARecord().x = 1

def test_cannot_assign_undeclared_fields():
    with pytest.raises(AttributeError):
        ARecord().set('z', 5)

def test_cannot_assign_wrong_type_to_fields():
    with pytest.raises(TypeError):
        ARecord().set('x', 'foo')

def test_cannot_construct_with_undeclared_fields():
    with pytest.raises(AttributeError):
        ARecord(z=5)

def test_cannot_construct_with_fields_of_wrong_type():
    with pytest.raises(TypeError):
        ARecord(x='foo')

def test_support_record_inheritance():
    class BRecord(ARecord):
        z = field()

    r = BRecord(x=1, y='foo', z='bar')

    assert isinstance(r, BRecord)
    assert isinstance(r, ARecord)
    assert r == {'x': 1, 'y': 'foo', 'z': 'bar'}

def test_single_type_spec():
    class A(PRecord):
        x = field(type=int)

    r = A(x=1)
    assert r.x == 1

    with pytest.raises(TypeError):
        r.set('x', 'foo')

def test_remove():
    r = ARecord(x=1, y='foo')
    r2 = r.remove('y')

    assert isinstance(r2, ARecord)
    assert r2 == {'x': 1}

def test_field_invariant_must_hold():
    class BRecord(PRecord):
        x = field(invariant=lambda x: (x > 1, 'x too small'))
        y = field(mandatory=True)

    try:
        BRecord(x=1)
        assert False
    except InvariantException as e:
        assert e.invariant_errors == ('x too small',)
        assert e.missing_fields == ('BRecord.y',)

def test_global_invariant_must_hold():
    class BRecord(PRecord):
        __invariant__ = lambda r: (r.x <= r.y, 'y smaller than x')
        x = field()
        y = field()

    BRecord(x=1, y=2)

    try:
        BRecord(x=2, y=1)
        assert False
    except InvariantException as e:
        assert e.invariant_errors == ('y smaller than x',)
        assert e.missing_fields == ()

def test_set_multiple_fields():
    a = ARecord(x=1, y='foo')
    b = a.set(x=2, y='bar')

    assert b == {'x': 2, 'y': 'bar'}

def test_initial_value():
    class BRecord(PRecord):
        x = field(initial=1)
        y = field(initial=2)

    a = BRecord()
    assert a.x == 1
    assert a.y == 2

def test_type_specification_must_be_a_type():
    with pytest.raises(TypeError):
        class BRecord(PRecord):
            x = field(type=1)

def test_initial_must_be_of_correct_type():
    with pytest.raises(TypeError):
        class BRecord(PRecord):
            x = field(type=int, initial='foo')

def test_invariant_must_be_callable():
    with pytest.raises(TypeError):
        class BRecord(PRecord):
            x = field(invariant='foo')

def test_global_invariants_are_inherited():
    class BRecord(PRecord):
        __invariant__ = lambda r: (r.x % r.y == 0, 'modulo')
        x = field()
        y = field()

    class CRecord(BRecord):
        __invariant__ = lambda r: (r.x > r.y, 'size')

    try:
        CRecord(x=5, y=3)
        assert False
    except InvariantException as e:
        assert e.invariant_errors == ('modulo',)

def test_global_invariants_must_be_callable():
    with pytest.raises(TypeError):
        class CRecord(PRecord):
            __invariant__ = 1

def test_repr():
    r = ARecord(x=1, y=2)
    assert repr(r) == 'ARecord(x=1, y=2)' or repr(r) == 'ARecord(y=2, x=1)'

def test_factory():
    class BRecord(PRecord):
        x = field(type=int, factory=int)

    assert BRecord(x=2.5) == {'x': 2}

def test_factory_must_be_callable():
    with pytest.raises(TypeError):
        class BRecord(PRecord):
            x = field(type=int, factory=1)

def test_nested_record_construction():
    class BRecord(PRecord):
        x = field(int, factory=int)

    class CRecord(PRecord):
        a = field()
        b = field(type=BRecord)

    r = CRecord.create({'a': 'foo', 'b': {'x': '5'}})
    assert isinstance(r, CRecord)
    assert isinstance(r.b, BRecord)
    assert r == {'a': 'foo', 'b': {'x': 5}}

def test_pickling():
    x = ARecord(x=2.0, y='bar')
    y = pickle.loads(pickle.dumps(x, -1))

    assert x == y
    assert isinstance(y, ARecord)

def test_all_invariant_errors_reported():
    class BRecord(PRecord):
        x = field(factory=int, invariant=lambda x: (x >= 0, 'x negative'))
        y = field(mandatory=True)

    class CRecord(PRecord):
        a = field(invariant=lambda x: (x != 0, 'a zero'))
        b = field(type=BRecord)

    try:
        CRecord.create({'a': 0, 'b': {'x': -5}})
        assert False
    except InvariantException as e:
        assert set(e.invariant_errors) == set(['x negative', 'a zero'])
        assert e.missing_fields == ('BRecord.y',)

def test_precord_factory_method_is_idempotent():
    class BRecord(PRecord):
        x = field()
        y = field()

    r = BRecord(x=1, y=2)
    assert BRecord.create(r) is r

def test_serialize():
    class BRecord(PRecord):
        d = field(type=datetime.date,
                  factory=lambda d: datetime.datetime.strptime(d, "%d%m%Y").date(),
                  serializer=lambda format, d: d.strftime('%Y-%m-%d') if format=='ISO' else d.strftime('%d%m%Y'))

    assert BRecord(d='14012015').serialize('ISO') == {'d': '2015-01-14'}
    assert BRecord(d='14012015').serialize('other') == {'d': '14012015'}

def test_nested_serialize():
    class BRecord(PRecord):
        d = field(serializer=lambda format, d: format)

    class CRecord(PRecord):
        b = field()

    serialized = CRecord(b=BRecord(d='foo')).serialize('bar')

    assert serialized == {'b': {'d': 'bar'}}
    assert isinstance(serialized, dict)

def test_serializer_must_be_callable():
    with pytest.raises(TypeError):
        class CRecord(PRecord):
            x = field(serializer=1)

def test_transform_without_update_returns_same_precord():
    r = ARecord(x=2.0, y='bar')
    assert r.transform([ny], lambda x: x) is r

class Application(PRecord):
    name = field(type=(six.text_type,) + six.string_types)
    image = field(type=(six.text_type,) + six.string_types)

class ApplicationVector(CheckedPVector):
    __type__ = Application

class Node(PRecord):
    applications = field(type=ApplicationVector)

def test_nested_create_serialize():
    node = Node(applications=[Application(name='myapp', image='myimage'),
                              Application(name='b', image='c')])

    node2 = Node.create({'applications': [{'name': 'myapp', 'image': 'myimage'},
                                          {'name': 'b', 'image': 'c'}]})

    assert node == node2

    serialized = node.serialize()
    restored = Node.create(serialized)

    assert restored == node
<|endoftext|>


<commit_before>#!/usr/bin/env python

import argparse
import json
import logging
import os
import shutil
import time
import re

from six.moves import configparser, input

from autotorrent.at import AutoTorrent
from autotorrent.clients import TORRENT_CLIENTS
from autotorrent.db import Database
from autotorrent.humanize import humanize_bytes

def query_yes_no(question, default="yes"):
    """Ask a yes/no question via raw_input() and return their answer.

    "question" is a string that is presented to the user.
    "default" is the presumed answer if the user just hits <Enter>.
        It must be "yes" (the default), "no" or None (meaning
        an answer is required of the user).

    The "answer" return value is True for "yes" or False for "no".
    """
    valid = {"yes": True, "y": True, "ye": True,
             "no": False, "n": False}
    if default is None:
        prompt = " [y/n] "
    elif default == "yes":
        prompt = " [Y/n] "
    elif default == "no":
        prompt = " [y/N] "
    else:
        raise ValueError("invalid default answer: '%s'" % default)

    while True:
        print(question + prompt)
        choice = input().lower()
        if default is not None and choice == '':
            return valid[default]
        elif choice in valid:
            return valid[choice]
        else:
            print("Please respond with 'yes' or 'no' (or 'y' or 'n').\n")

def commandline_handler():
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--config", dest="config_file", default="autotorrent.conf", help="Path to config file")
    parser.add_argument("-l", "--client", dest="client", default="default", help="Name of client to use (when multiple configured)")

    parser.add_argument("--create_config", dest="create_config_file", nargs='?', const='autotorrent.conf', default=None, help="Creates a new configuration file")
    parser.add_argument("-t", "--test_connection", action="store_true", dest="test_connection", default=False, help='Tests the connection to the torrent client')
    parser.add_argument("--dry-run", nargs='?', const='txt', default=None, dest="dry_run", choices=['txt', 'json'], help="Don't do any actual adding, just scan for files needed for torrents.")
    parser.add_argument("-r", "--rebuild", dest="rebuild", default=False, help='Rebuild the database', nargs='*')
    parser.add_argument("-a", "--addfile", dest="addfile", default=False, help='Add a new torrent file to client', nargs='+')
    parser.add_argument("-d", "--delete_torrents", action="store_true", dest="delete_torrents", default=False, help='Delete torrents when they are added to the client')
    parser.add_argument("--verbose", help="increase output verbosity", action="store_true", dest="verbose")
    parser.add_argument("-o", "--loopmode", dest="loopmode", default=False, help='Enable loop mode (scan directory '
                                                                                'for torrents every few seconds)', nargs='?')

    args = parser.parse_args()

    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.ERROR)

    if args.create_config_file: # autotorrent.conf
        if os.path.exists(args.create_config_file):
            parser.error("Target %r already exists, not creating" % args.create_config_file)
        else:
            src = os.path.join(os.path.dirname(__file__), 'autotorrent.conf.dist')
            shutil.copy(src, args.create_config_file)
            print('Created configuration file %r' % args.create_config_file)

        client_config = {
            'client': 'rtorrent',
            'url': 'http://user:pass@127.0.0.1/RPC2',
            'label': 'autotorrent',
        }

        if query_yes_no('Do you want to try and auto-configure torrent client?'):
            working_clients = []
            for client_name, cls in TORRENT_CLIENTS.items():
                obj = cls.auto_config()
                try:
                    if obj.test_connection():
                        working_clients.append(obj)
                except:
                    continue

            if working_clients:
                print('Found %i clients - please choose a client to use' % len(working_clients))
                for i, client in enumerate(working_clients, 1):
                    print('[%i] %s' % (i, client.identifier))
                print('[0] None of the above - do not auto-configure any client\n')

                while True:
                    error = False
                    try:
                        choice = int(input('> '))
                    except ValueError:
                        error = True
                    else:
                        if len(working_clients) < choice or choice < 0:
                            error = True

                    if error:
                        print('Invalid choice, please choose again')
                    else:
                        if choice > 0:
                            client = working_clients[choice-1]
                            print('Setting client to %s' % client.identifier)
                            client_config = client.get_config()
                            client_config['client'] = client.identifier

                        break
            else:
                print('Unable to auto-detect any clients, you will have to configure it manually.')

        config = configparser.ConfigParser()
        config.read(args.create_config_file)
        for k, v in client_config.items():
            config.set('client', k, v)

        with open(args.create_config_file, 'w') as configfile:
            config.write(configfile)

        quit()

    if not os.path.isfile(args.config_file):
        parser.error("Config file not found %r" % args.config_file)

    print('Using config file %s' % args.config_file)
    config = configparser.ConfigParser()
    config.read(args.config_file)

    current_path = os.getcwd()
    os.chdir(os.path.dirname(os.path.abspath(args.config_file))) # Changing directory to where the config file is.

    if not config.has_section('general'):
        parser.error('AutoTorrent is not properly configured, please edit %r' % args.config_file)
        quit(1)

    i = 1
    disks = []
    while config.has_option('disks', 'disk%s' % i):
        disks.append(config.get('disks', 'disk%s' % i))
        i += 1

    scan_mode = set(config.get('general', 'scan_mode').split(','))

    exact_mode = 'exact' in scan_mode
    unsplitable_mode = 'unsplitable' in scan_mode
    normal_mode = 'normal' in scan_mode

    hash_name_mode = 'hash_name' in scan_mode
    hash_size_mode = 'hash_size' in scan_mode
    hash_slow_mode = 'hash_slow' in scan_mode

    db = Database(config.get('general', 'db'), disks,
                  config.get('general', 'ignore_files').split(','),
                  normal_mode, unsplitable_mode, exact_mode,
                  hash_name_mode, hash_size_mode, hash_slow_mode)

    client_option = 'client'
    if args.client != 'default':
        client_option += '-%s' % args.client

    try:
        client_name = config.get(client_option, 'client')
    except configparser.NoSectionError:
        print('It seems like %r is not a configured client' % args.client)
        quit(1)

    if client_name not in TORRENT_CLIENTS:
        print('Unknown client %r - Known clients are: %s' % (client_name, ', '.join(TORRENT_CLIENTS.keys())))
        quit(1)

    client_options = dict(config.items(client_option))
    client_options.pop('client')
    client = TORRENT_CLIENTS[client_name](**client_options)

    at = AutoTorrent(
        db,
        client,
        config.get('general', 'store_path'),
        config.getint('general', 'add_limit_size'),
        config.getfloat('general', 'add_limit_percent'),
        args.delete_torrents,
        (config.get('general', 'link_type') if config.has_option('general', 'link_type') else 'soft'),
    )

    if args.test_connection:
        proxy_test_result = client.test_connection()
        if proxy_test_result:
            print('Connected to torrent client successfully!')
            print('  result: %s' % proxy_test_result)

    if isinstance(args.rebuild, list):
        if args.rebuild:
            print('Adding new folders to database')
            db.rebuild(args.rebuild)
            print('Added to database')
        else:
            print('Rebuilding database')
            db.rebuild()
            print('Database rebuilt')

    if args.addfile:
        addtfile(at, current_path, args.addfile, args.dry_run)

    if args.loopmode:
        print(config.get('general', 'store_path'))
        print('Entering loop mode (%s) (ctrl-c to exit)' % args.loopmode)
        while 1 == 1:
            for fn in os.listdir(args.loopmode):
                if fn.endswith('.torrent'):
                    fn_woext = os.path.splitext(fn)[0]
                    fn_scenename = re.search('-(.*)$', fn_woext).group(1).replace(' ', '.').lower()
                    print('!FOUND %s (%s)' % (fn_woext, fn_scenename))

                    at.populate_torrents_seeded_names()
                    print('!  There is currently %i seeded torrents in client' % len(at.torrents_seeded_names))

                    # Check if torrent exists
                    added = False
                    for thash, tname in at.torrents_seeded_names:
                        if tname == fn_scenename:
                            # If exists, check if seeding
                            seeding = at.get_complete(thash)
                            print('!  This release is already in the client and is %s' % ("seeding" if seeding else
                                                                                "downloading"))
                            # If seeding
                            if seeding:
                                # Add to cross-seed
                                print('!  Adding torrent in cross-seed mode')
                                addtfile(at, args.loopmode, [fn], args.dry_run)

                                # delete torrent file
                                print('!  Deleting file')
                                os.remove(os.path.join(args.loopmode, fn))

                                # rebuild files db
                                print('!  Adding new folders to database')
                                db.rebuild([config.get('general', 'store_path')])
                                added = True
                                break
                            else:
                                print('!  Skipping')
                                break
                    if not added:
                        # If not exists, add new
                        print("!  Adding torrent")

                        # delete torrent file
                        print("!  Deleting file")
                        os.remove(os.path.join(current_path, fn))

                        # rebuild files db
                        print('!  Adding new folders to database')
                        db.rebuild([config.get('general', 'store_path')])

                    print('Entering loop mode (%s) (ctrl-c to exit)' % args.loopmode)

            time.sleep(5)

def addtfile(at, current_path, afiles, adry_run):
    dry_run = bool(adry_run)
    dry_run_data = []
    if not dry_run:
        at.populate_torrents_seeded()

    for torrent in afiles:
        result = at.handle_torrentfile(os.path.join(current_path, torrent), dry_run)
        if dry_run:
            dry_run_data.append({
                'torrent': torrent,
                'found_bytes': result[0],
                'missing_bytes': result[1],
                'would_add': not result[2],
                'local_files': result[3],
            })

    if dry_run:
        if dry_run == 'json':
            print(json.dumps(dry_run_data))
        elif dry_run == 'txt':
            for torrent in dry_run_data:
                print('Torrent: %s' % torrent['torrent'])
                print(' Found data: %s - Missing data: %s - Would add: %s' % (humanize_bytes(torrent['found_bytes']),
                                                                              humanize_bytes(torrent['missing_bytes']),
                                                                              torrent['would_add'] and 'Yes' or 'No'))
                print(' Local files used:')
                for f in torrent['local_files']:
                    print('  %s' % f)
                print('')

if __name__ == '__main__':
    commandline_handler()
<commit_msg>Starting loop mode coding<commit_after>#!/usr/bin/env python

import argparse
import json
import logging
import os
import shutil
import time
import re

from six.moves import configparser, input

from autotorrent.at import AutoTorrent
from autotorrent.clients import TORRENT_CLIENTS
from autotorrent.db import Database
from autotorrent.humanize import humanize_bytes

def query_yes_no(question, default="yes"):
    """Ask a yes/no question via raw_input() and return their answer.

    "question" is a string that is presented to the user.
    "default" is the presumed answer if the user just hits <Enter>.
        It must be "yes" (the default), "no" or None (meaning
        an answer is required of the user).

    The "answer" return value is True for "yes" or False for "no".
    """
    valid = {"yes": True, "y": True, "ye": True,
             "no": False, "n": False}
    if default is None:
        prompt = " [y/n] "
    elif default == "yes":
        prompt = " [Y/n] "
    elif default == "no":
        prompt = " [y/N] "
    else:
        raise ValueError("invalid default answer: '%s'" % default)

    while True:
        print(question + prompt)
        choice = input().lower()
        if default is not None and choice == '':
            return valid[default]
        elif choice in valid:
            return valid[choice]
        else:
            print("Please respond with 'yes' or 'no' (or 'y' or 'n').\n")

def commandline_handler():
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--config", dest="config_file", default="autotorrent.conf", help="Path to config file")
    parser.add_argument("-l", "--client", dest="client", default="default", help="Name of client to use (when multiple configured)")

    parser.add_argument("--create_config", dest="create_config_file", nargs='?', const='autotorrent.conf', default=None, help="Creates a new configuration file")
    parser.add_argument("-t", "--test_connection", action="store_true", dest="test_connection", default=False, help='Tests the connection to the torrent client')
    parser.add_argument("--dry-run", nargs='?', const='txt', default=None, dest="dry_run", choices=['txt', 'json'], help="Don't do any actual adding, just scan for files needed for torrents.")
    parser.add_argument("-r", "--rebuild", dest="rebuild", default=False, help='Rebuild the database', nargs='*')
    parser.add_argument("-a", "--addfile", dest="addfile", default=False, help='Add a new torrent file to client', nargs='+')
    parser.add_argument("-d", "--delete_torrents", action="store_true", dest="delete_torrents", default=False, help='Delete torrents when they are added to the client')
    parser.add_argument("--verbose", help="increase output verbosity", action="store_true", dest="verbose")
    parser.add_argument("-o", "--loopmode", dest="loopmode", default=False, help='Enable loop mode (scan directory '
                                                                                'for torrents every few seconds)', nargs='?')

    args = parser.parse_args()

    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.ERROR)

    if args.create_config_file: # autotorrent.conf
        if os.path.exists(args.create_config_file):
            parser.error("Target %r already exists, not creating" % args.create_config_file)
        else:
            src = os.path.join(os.path.dirname(__file__), 'autotorrent.conf.dist')
            shutil.copy(src, args.create_config_file)
            print('Created configuration file %r' % args.create_config_file)

        client_config = {
            'client': 'rtorrent',
            'url': 'http://user:pass@127.0.0.1/RPC2',
            'label': 'autotorrent',
        }

        if query_yes_no('Do you want to try and auto-configure torrent client?'):
            working_clients = []
            for client_name, cls in TORRENT_CLIENTS.items():
                obj = cls.auto_config()
                try:
                    if obj.test_connection():
                        working_clients.append(obj)
                except:
                    continue

            if working_clients:
                print('Found %i clients - please choose a client to use' % len(working_clients))
                for i, client in enumerate(working_clients, 1):
                    print('[%i] %s' % (i, client.identifier))
                print('[0] None of the above - do not auto-configure any client\n')

                while True:
                    error = False
                    try:
                        choice = int(input('> '))
                    except ValueError:
                        error = True
                    else:
                        if len(working_clients) < choice or choice < 0:
                            error = True

                    if error:
                        print('Invalid choice, please choose again')
                    else:
                        if choice > 0:
                            client = working_clients[choice-1]
                            print('Setting client to %s' % client.identifier)
                            client_config = client.get_config()
                            client_config['client'] = client.identifier

                        break
            else:
                print('Unable to auto-detect any clients, you will have to configure it manually.')

        config = configparser.ConfigParser()
        config.read(args.create_config_file)
        for k, v in client_config.items():
            config.set('client', k, v)

        with open(args.create_config_file, 'w') as configfile:
            config.write(configfile)

        quit()

    if not os.path.isfile(args.config_file):
        parser.error("Config file not found %r" % args.config_file)

    print('Using config file %s' % args.config_file)
    config = configparser.ConfigParser()
    config.read(args.config_file)

    current_path = os.getcwd()
    os.chdir(os.path.dirname(os.path.abspath(args.config_file))) # Changing directory to where the config file is.

    if not config.has_section('general'):
        parser.error('AutoTorrent is not properly configured, please edit %r' % args.config_file)
        quit(1)

    i = 1
    disks = []
    while config.has_option('disks', 'disk%s' % i):
        disks.append(config.get('disks', 'disk%s' % i))
        i += 1

    scan_mode = set(config.get('general', 'scan_mode').split(','))

    exact_mode = 'exact' in scan_mode
    unsplitable_mode = 'unsplitable' in scan_mode
    normal_mode = 'normal' in scan_mode

    hash_name_mode = 'hash_name' in scan_mode
    hash_size_mode = 'hash_size' in scan_mode
    hash_slow_mode = 'hash_slow' in scan_mode

    db = Database(config.get('general', 'db'), disks,
                  config.get('general', 'ignore_files').split(','),
                  normal_mode, unsplitable_mode, exact_mode,
                  hash_name_mode, hash_size_mode, hash_slow_mode)

    client_option = 'client'
    if args.client != 'default':
        client_option += '-%s' % args.client

    try:
        client_name = config.get(client_option, 'client')
    except configparser.NoSectionError:
        print('It seems like %r is not a configured client' % args.client)
        quit(1)

    if client_name not in TORRENT_CLIENTS:
        print('Unknown client %r - Known clients are: %s' % (client_name, ', '.join(TORRENT_CLIENTS.keys())))
        quit(1)

    client_options = dict(config.items(client_option))
    client_options.pop('client')
    client = TORRENT_CLIENTS[client_name](**client_options)

    at = AutoTorrent(
        db,
        client,
        config.get('general', 'store_path'),
        config.getint('general', 'add_limit_size'),
        config.getfloat('general', 'add_limit_percent'),
        args.delete_torrents,
        (config.get('general', 'link_type') if config.has_option('general', 'link_type') else 'soft'),
    )

    if args.test_connection:
        proxy_test_result = client.test_connection()
        if proxy_test_result:
            print('Connected to torrent client successfully!')
            print('  result: %s' % proxy_test_result)

    if isinstance(args.rebuild, list):
        if args.rebuild:
            print('Adding new folders to database')
            db.rebuild(args.rebuild)
            print('Added to database')
        else:
            print('Rebuilding database')
            db.rebuild()
            print('Database rebuilt')

    if args.addfile:
        addtfile(at, current_path, args.addfile, args.dry_run)

    if args.loopmode:
        print(config.get('general', 'store_path'))
        print('Entering loop mode (%s) (ctrl-c to exit)' % args.loopmode)
        while 1 == 1:
            for fn in os.listdir(args.loopmode):
                if fn.endswith('.torrent'):
                    fn_woext = os.path.splitext(fn)[0]
                    fn_scenename = re.search('-(.*)$', fn_woext).group(1).replace(' ', '.').lower()
                    print('!FOUND %s (%s)' % (fn_woext, fn_scenename))

                    at.populate_torrents_seeded_names()
                    print('!  There is currently %i torrents in client' % len(at.torrents_seeded_names))

                    # Check if torrent exists
                    added = False
                    for thash, tname in at.torrents_seeded_names:
                        if tname == fn_scenename:
                            # If exists, check if seeding
                            seeding = at.get_complete(thash)
                            print('!  This release is already in the client and is %s' % ("seeding" if seeding else
                                                                                "downloading"))
                            # If seeding
                            if seeding:
                                # Add to cross-seed
                                print('!  Adding torrent in cross-seed mode')
                                addtfile(at, args.loopmode, [fn], args.dry_run)

                                # delete torrent file
                                print('!  Deleting file')
                                os.remove(os.path.join(args.loopmode, fn))

                                # rebuild files db
                                print('!  Adding new folders to database')
                                db.rebuild([config.get('general', 'store_path')])
                                added = True
                                break
                            else:
                                print('!  Skipping')
                                break
                    if not added:
                        # If not exists, add new
                        print("!  Adding torrent")

                        # delete torrent file
                        print("!  Deleting file")
                        os.remove(os.path.join(current_path, fn))

                        # rebuild files db
                        print('!  Adding new folders to database')
                        db.rebuild([config.get('general', 'store_path')])

                    print('\n\nEntering loop mode (%s) (ctrl-c to exit)' % args.loopmode)

            time.sleep(5)

def addtfile(at, current_path, afiles, adry_run):
    dry_run = bool(adry_run)
    dry_run_data = []
    if not dry_run:
        at.populate_torrents_seeded()

    for torrent in afiles:
        result = at.handle_torrentfile(os.path.join(current_path, torrent), dry_run)
        if dry_run:
            dry_run_data.append({
                'torrent': torrent,
                'found_bytes': result[0],
                'missing_bytes': result[1],
                'would_add': not result[2],
                'local_files': result[3],
            })

    if dry_run:
        if dry_run == 'json':
            print(json.dumps(dry_run_data))
        elif dry_run == 'txt':
            for torrent in dry_run_data:
                print('Torrent: %s' % torrent['torrent'])
                print(' Found data: %s - Missing data: %s - Would add: %s' % (humanize_bytes(torrent['found_bytes']),
                                                                              humanize_bytes(torrent['missing_bytes']),
                                                                              torrent['would_add'] and 'Yes' or 'No'))
                print(' Local files used:')
                for f in torrent['local_files']:
                    print('  %s' % f)
                print('')

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


<commit_before>#!/usr/bin/python
#
# Copyright 2013, Google Inc. All rights reserved.
# Use of this source code is governed by a BSD-style license that can
# be found in the LICENSE file.

import base64
import logging
import threading
import struct
import time
import unittest

from vtdb import keyrange_constants
from vtdb import keyspace

import environment
import utils
import tablet

from zk import zkocc

SHARDED_KEYSPACE = "TEST_KEYSPACE_SHARDED"
UNSHARDED_KEYSPACE = "TEST_KEYSPACE_UNSHARDED"

# shards for SHARDED_KEYSPACE
# range "" - 80
shard_0_master = tablet.Tablet()
shard_0_replica = tablet.Tablet()
shard_0_rdonly = tablet.Tablet()
# range 80 - ""
shard_1_master = tablet.Tablet()
shard_1_replica = tablet.Tablet()
shard_1_rdonly = tablet.Tablet()

# shard for UNSHARDED_KEYSPACE
unsharded_master = tablet.Tablet()
unsharded_replica = tablet.Tablet()
unsharded_rdonly = tablet.Tablet()

vtgate_server = None
vtgate_port = None

shard_names = ['-80', '80-']
shard_kid_map = {'-80': [527875958493693904, 626750931627689502,
                         345387386794260318, 332484755310826578,
                         1842642426274125671, 1326307661227634652,
                         1761124146422844620, 1661669973250483744,
                         3361397649937244239, 2444880764308344533],
                 '80-': [9767889778372766922, 9742070682920810358,
                         10296850775085416642, 9537430901666854108,
                         10440455099304929791, 11454183276974683945,
                         11185910247776122031, 10460396697869122981,
                         13379616110062597001, 12826553979133932576],
                 }

create_vt_insert_test = '''create table vt_insert_test (
id bigint auto_increment,
msg varchar(64),
keyspace_id bigint(20) unsigned NOT NULL,
primary key (id)
) Engine=InnoDB'''

def setUpModule():
  try:
    environment.topo_server_setup()

    setup_procs = [
        shard_0_master.init_mysql(),
        shard_0_replica.init_mysql(),
        shard_0_rdonly.init_mysql(),
        shard_1_master.init_mysql(),
        shard_1_replica.init_mysql(),
        shard_1_rdonly.init_mysql(),
        unsharded_master.init_mysql(),
        unsharded_replica.init_mysql(),
        unsharded_rdonly.init_mysql(),
        ]
    utils.wait_procs(setup_procs)
    setup_tablets()
  except:
    tearDownModule()
    raise

def tearDownModule():
  if utils.options.skip_teardown:
    return

  global vtgate_server
  utils.vtgate_kill(vtgate_server)
  tablet.kill_tablets([shard_0_master, shard_0_replica, shard_0_rdonly,
                      shard_1_master, shard_1_replica, shard_1_rdonly])
  teardown_procs = [
      shard_0_master.teardown_mysql(),
      shard_0_replica.teardown_mysql(),
      shard_0_rdonly.teardown_mysql(),
      shard_1_master.teardown_mysql(),
      shard_1_replica.teardown_mysql(),
      shard_1_rdonly.teardown_mysql(),
      unsharded_master.teardown_mysql(),
      unsharded_replica.teardown_mysql(),
      unsharded_rdonly.teardown_mysql(),
      ]
  utils.wait_procs(teardown_procs, raise_on_error=False)

  environment.topo_server_teardown()
  utils.kill_sub_processes()
  utils.remove_tmp_files()

  shard_0_master.remove_tree()
  shard_0_replica.remove_tree()
  shard_0_rdonly.remove_tree()
  shard_1_master.remove_tree()
  shard_1_replica.remove_tree()
  shard_1_rdonly.remove_tree()
  shard_1_rdonly.remove_tree()
  shard_1_rdonly.remove_tree()
  unsharded_master.remove_tree()
  unsharded_replica.remove_tree()
  unsharded_rdonly.remove_tree()

def setup_tablets():
  global vtgate_server
  global vtgate_port

  setup_sharded_keyspace()
  setup_unsharded_keyspace()
  vtgate_server, vtgate_port = utils.vtgate_start()

def setup_sharded_keyspace():
  utils.run_vtctl(['CreateKeyspace', SHARDED_KEYSPACE])
  utils.run_vtctl(['SetKeyspaceShardingInfo', '-force', SHARDED_KEYSPACE,
                   'keyspace_id', 'uint64'])
  shard_0_master.init_tablet('master', keyspace=SHARDED_KEYSPACE, shard='-80')
  shard_0_replica.init_tablet('replica', keyspace=SHARDED_KEYSPACE, shard='-80')
  shard_0_rdonly.init_tablet('rdonly', keyspace=SHARDED_KEYSPACE, shard='-80')
  shard_1_master.init_tablet('master', keyspace=SHARDED_KEYSPACE, shard='80-')
  shard_1_replica.init_tablet('replica', keyspace=SHARDED_KEYSPACE,  shard='80-')
  shard_1_rdonly.init_tablet('rdonly', keyspace=SHARDED_KEYSPACE,  shard='80-')

  utils.run_vtctl(['RebuildKeyspaceGraph', SHARDED_KEYSPACE,], auto_log=True)

  for t in [shard_0_master, shard_0_replica, shard_0_rdonly, shard_1_master, shard_1_replica, shard_1_rdonly]:
    t.create_db('vt_test_keyspace_sharded')
    t.mquery(shard_0_master.dbname, create_vt_insert_test)
    t.start_vttablet(wait_for_state=None)

  for t in [shard_0_master, shard_0_replica, shard_1_master, shard_1_replica]:
    t.wait_for_vttablet_state('SERVING')

  utils.run_vtctl(['ReparentShard', '-force', '%s/-80' % SHARDED_KEYSPACE,
                   shard_0_master.tablet_alias], auto_log=True)
  utils.run_vtctl(['ReparentShard', '-force', '%s/80-' % SHARDED_KEYSPACE,
                   shard_1_master.tablet_alias], auto_log=True)

  utils.run_vtctl(['RebuildKeyspaceGraph', SHARDED_KEYSPACE],
                   auto_log=True)

  utils.check_srv_keyspace('test_nj', SHARDED_KEYSPACE,
                           'Partitions(master): -80 80-\n' +
                           'Partitions(rdonly): -80 80-\n' +
                           'Partitions(replica): -80 80-\n' +
                           'TabletTypes: master,rdonly,replica')

def setup_unsharded_keyspace():
  utils.run_vtctl(['CreateKeyspace', UNSHARDED_KEYSPACE])
  utils.run_vtctl(['SetKeyspaceShardingInfo', '-force', UNSHARDED_KEYSPACE,
                   'keyspace_id', 'uint64'])
  unsharded_master.init_tablet('master', keyspace=UNSHARDED_KEYSPACE, shard='0')
  unsharded_replica.init_tablet('replica', keyspace=UNSHARDED_KEYSPACE, shard='0')
  unsharded_rdonly.init_tablet('rdonly', keyspace=UNSHARDED_KEYSPACE, shard='0')

  utils.run_vtctl(['RebuildKeyspaceGraph', UNSHARDED_KEYSPACE,], auto_log=True)

  for t in [unsharded_master, unsharded_replica, unsharded_rdonly]:
    t.create_db('vt_test_keyspace_unsharded')
    t.mquery(unsharded_master.dbname, create_vt_insert_test)
    t.start_vttablet(wait_for_state=None)

  for t in [unsharded_master, unsharded_replica, unsharded_rdonly]:
    t.wait_for_vttablet_state('SERVING')

  utils.run_vtctl(['ReparentShard', '-force', '%s/0' % UNSHARDED_KEYSPACE,
                   unsharded_master.tablet_alias], auto_log=True)

  utils.run_vtctl(['RebuildKeyspaceGraph', UNSHARDED_KEYSPACE],
                   auto_log=True)

  utils.check_srv_keyspace('test_nj', UNSHARDED_KEYSPACE,
                           'Partitions(master): -\n' +
                           'Partitions(rdonly): -\n' +
                           'Partitions(replica): -\n' +
                           'TabletTypes: master,rdonly,replica')

ALL_DB_TYPES = ['master', 'replica', 'rdonly']

class TestKeyspace(unittest.TestCase):
  def _read_keyspace(self, keyspace_name):
    global vtgate_port
    vtgate_client = zkocc.ZkOccConnection("localhost:%u" % vtgate_port,
                                        "test_nj", 30.0)
    return keyspace.read_keyspace(vtgate_client, keyspace_name)

  def test_shard_count(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    self.assertEqual(sharded_ks.shard_count, 2)
    for db_type in ALL_DB_TYPES:
      self.assertEqual(sharded_ks.get_shard_count(db_type), 2)
    unsharded_ks = self._read_keyspace(UNSHARDED_KEYSPACE)
    self.assertEqual(unsharded_ks.shard_count, 1)
    for db_type in ALL_DB_TYPES:
      self.assertEqual(unsharded_ks.get_shard_count(db_type), 1)

  def test_shard_names(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    self.assertEqual(sharded_ks.shard_names, ['-80', '80-'])
    for db_type in ALL_DB_TYPES:
      self.assertEqual(sharded_ks.get_shard_names(db_type), ['-80', '80-'])
    unsharded_ks = self._read_keyspace(UNSHARDED_KEYSPACE)
    self.assertEqual(unsharded_ks.shard_names, ['0'])
    for db_type in ALL_DB_TYPES:
      self.assertEqual(unsharded_ks.get_shard_names(db_type), ['0'])

  def test_shard_max_keys(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    want = ['80', '']
    for i, smk in enumerate(sharded_ks.shard_max_keys):
      self.assertEqual(smk.encode('hex').upper(), want[i])
    for db_type in ALL_DB_TYPES:
      for i, smk in enumerate(sharded_ks.get_shard_max_keys(db_type)):
        self.assertEqual(smk.encode('hex').upper(), want[i])
    unsharded_ks = self._read_keyspace(UNSHARDED_KEYSPACE)
    self.assertEqual(unsharded_ks.shard_max_keys, None)
    for db_type in ALL_DB_TYPES:
      self.assertEqual(unsharded_ks.get_shard_max_keys(db_type), [''])

  def test_db_types(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    self.assertEqual(set(sharded_ks.db_types), set(ALL_DB_TYPES))
    unsharded_ks = self._read_keyspace(UNSHARDED_KEYSPACE)
    self.assertEqual(set(unsharded_ks.db_types), set(ALL_DB_TYPES))

  def test_keyspace_id_to_shard_index(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    for i, sn in enumerate(shard_names):
      for keyspace_id in shard_kid_map[sn]:
        self.assertEqual(sharded_ks.keyspace_id_to_shard_index(keyspace_id), i)
        self.assertEqual(sharded_ks.keyspace_id_to_shard_index_for_db_type(keyspace_id, 'master'), i)

  def test_keyspace_id_to_shard_index(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    for _, sn in enumerate(shard_names):
      for keyspace_id in shard_kid_map[sn]:
        self.assertEqual(sharded_ks.keyspace_id_to_shard_name_for_db_type(keyspace_id, 'master'), sn)
    unsharded_ks = self._read_keyspace(UNSHARDED_KEYSPACE)
    for keyspace_id in shard_kid_map[sn]:
      self.assertEqual(unsharded_ks.keyspace_id_to_shard_name_for_db_type(keyspace_id, 'master'), '0')

if __name__ == '__main__':
  utils.main()
<commit_msg>Minor fix to the test name.<commit_after>#!/usr/bin/python
#
# Copyright 2013, Google Inc. All rights reserved.
# Use of this source code is governed by a BSD-style license that can
# be found in the LICENSE file.

import base64
import logging
import threading
import struct
import time
import unittest

from vtdb import keyrange_constants
from vtdb import keyspace

import environment
import utils
import tablet

from zk import zkocc

SHARDED_KEYSPACE = "TEST_KEYSPACE_SHARDED"
UNSHARDED_KEYSPACE = "TEST_KEYSPACE_UNSHARDED"

# shards for SHARDED_KEYSPACE
# range "" - 80
shard_0_master = tablet.Tablet()
shard_0_replica = tablet.Tablet()
shard_0_rdonly = tablet.Tablet()
# range 80 - ""
shard_1_master = tablet.Tablet()
shard_1_replica = tablet.Tablet()
shard_1_rdonly = tablet.Tablet()

# shard for UNSHARDED_KEYSPACE
unsharded_master = tablet.Tablet()
unsharded_replica = tablet.Tablet()
unsharded_rdonly = tablet.Tablet()

vtgate_server = None
vtgate_port = None

shard_names = ['-80', '80-']
shard_kid_map = {'-80': [527875958493693904, 626750931627689502,
                         345387386794260318, 332484755310826578,
                         1842642426274125671, 1326307661227634652,
                         1761124146422844620, 1661669973250483744,
                         3361397649937244239, 2444880764308344533],
                 '80-': [9767889778372766922, 9742070682920810358,
                         10296850775085416642, 9537430901666854108,
                         10440455099304929791, 11454183276974683945,
                         11185910247776122031, 10460396697869122981,
                         13379616110062597001, 12826553979133932576],
                 }

create_vt_insert_test = '''create table vt_insert_test (
id bigint auto_increment,
msg varchar(64),
keyspace_id bigint(20) unsigned NOT NULL,
primary key (id)
) Engine=InnoDB'''

def setUpModule():
  try:
    environment.topo_server_setup()

    setup_procs = [
        shard_0_master.init_mysql(),
        shard_0_replica.init_mysql(),
        shard_0_rdonly.init_mysql(),
        shard_1_master.init_mysql(),
        shard_1_replica.init_mysql(),
        shard_1_rdonly.init_mysql(),
        unsharded_master.init_mysql(),
        unsharded_replica.init_mysql(),
        unsharded_rdonly.init_mysql(),
        ]
    utils.wait_procs(setup_procs)
    setup_tablets()
  except:
    tearDownModule()
    raise

def tearDownModule():
  if utils.options.skip_teardown:
    return

  global vtgate_server
  utils.vtgate_kill(vtgate_server)
  tablet.kill_tablets([shard_0_master, shard_0_replica, shard_0_rdonly,
                      shard_1_master, shard_1_replica, shard_1_rdonly])
  teardown_procs = [
      shard_0_master.teardown_mysql(),
      shard_0_replica.teardown_mysql(),
      shard_0_rdonly.teardown_mysql(),
      shard_1_master.teardown_mysql(),
      shard_1_replica.teardown_mysql(),
      shard_1_rdonly.teardown_mysql(),
      unsharded_master.teardown_mysql(),
      unsharded_replica.teardown_mysql(),
      unsharded_rdonly.teardown_mysql(),
      ]
  utils.wait_procs(teardown_procs, raise_on_error=False)

  environment.topo_server_teardown()
  utils.kill_sub_processes()
  utils.remove_tmp_files()

  shard_0_master.remove_tree()
  shard_0_replica.remove_tree()
  shard_0_rdonly.remove_tree()
  shard_1_master.remove_tree()
  shard_1_replica.remove_tree()
  shard_1_rdonly.remove_tree()
  shard_1_rdonly.remove_tree()
  shard_1_rdonly.remove_tree()
  unsharded_master.remove_tree()
  unsharded_replica.remove_tree()
  unsharded_rdonly.remove_tree()

def setup_tablets():
  global vtgate_server
  global vtgate_port

  setup_sharded_keyspace()
  setup_unsharded_keyspace()
  vtgate_server, vtgate_port = utils.vtgate_start()

def setup_sharded_keyspace():
  utils.run_vtctl(['CreateKeyspace', SHARDED_KEYSPACE])
  utils.run_vtctl(['SetKeyspaceShardingInfo', '-force', SHARDED_KEYSPACE,
                   'keyspace_id', 'uint64'])
  shard_0_master.init_tablet('master', keyspace=SHARDED_KEYSPACE, shard='-80')
  shard_0_replica.init_tablet('replica', keyspace=SHARDED_KEYSPACE, shard='-80')
  shard_0_rdonly.init_tablet('rdonly', keyspace=SHARDED_KEYSPACE, shard='-80')
  shard_1_master.init_tablet('master', keyspace=SHARDED_KEYSPACE, shard='80-')
  shard_1_replica.init_tablet('replica', keyspace=SHARDED_KEYSPACE,  shard='80-')
  shard_1_rdonly.init_tablet('rdonly', keyspace=SHARDED_KEYSPACE,  shard='80-')

  utils.run_vtctl(['RebuildKeyspaceGraph', SHARDED_KEYSPACE,], auto_log=True)

  for t in [shard_0_master, shard_0_replica, shard_0_rdonly, shard_1_master, shard_1_replica, shard_1_rdonly]:
    t.create_db('vt_test_keyspace_sharded')
    t.mquery(shard_0_master.dbname, create_vt_insert_test)
    t.start_vttablet(wait_for_state=None)

  for t in [shard_0_master, shard_0_replica, shard_1_master, shard_1_replica]:
    t.wait_for_vttablet_state('SERVING')

  utils.run_vtctl(['ReparentShard', '-force', '%s/-80' % SHARDED_KEYSPACE,
                   shard_0_master.tablet_alias], auto_log=True)
  utils.run_vtctl(['ReparentShard', '-force', '%s/80-' % SHARDED_KEYSPACE,
                   shard_1_master.tablet_alias], auto_log=True)

  utils.run_vtctl(['RebuildKeyspaceGraph', SHARDED_KEYSPACE],
                   auto_log=True)

  utils.check_srv_keyspace('test_nj', SHARDED_KEYSPACE,
                           'Partitions(master): -80 80-\n' +
                           'Partitions(rdonly): -80 80-\n' +
                           'Partitions(replica): -80 80-\n' +
                           'TabletTypes: master,rdonly,replica')

def setup_unsharded_keyspace():
  utils.run_vtctl(['CreateKeyspace', UNSHARDED_KEYSPACE])
  utils.run_vtctl(['SetKeyspaceShardingInfo', '-force', UNSHARDED_KEYSPACE,
                   'keyspace_id', 'uint64'])
  unsharded_master.init_tablet('master', keyspace=UNSHARDED_KEYSPACE, shard='0')
  unsharded_replica.init_tablet('replica', keyspace=UNSHARDED_KEYSPACE, shard='0')
  unsharded_rdonly.init_tablet('rdonly', keyspace=UNSHARDED_KEYSPACE, shard='0')

  utils.run_vtctl(['RebuildKeyspaceGraph', UNSHARDED_KEYSPACE,], auto_log=True)

  for t in [unsharded_master, unsharded_replica, unsharded_rdonly]:
    t.create_db('vt_test_keyspace_unsharded')
    t.mquery(unsharded_master.dbname, create_vt_insert_test)
    t.start_vttablet(wait_for_state=None)

  for t in [unsharded_master, unsharded_replica, unsharded_rdonly]:
    t.wait_for_vttablet_state('SERVING')

  utils.run_vtctl(['ReparentShard', '-force', '%s/0' % UNSHARDED_KEYSPACE,
                   unsharded_master.tablet_alias], auto_log=True)

  utils.run_vtctl(['RebuildKeyspaceGraph', UNSHARDED_KEYSPACE],
                   auto_log=True)

  utils.check_srv_keyspace('test_nj', UNSHARDED_KEYSPACE,
                           'Partitions(master): -\n' +
                           'Partitions(rdonly): -\n' +
                           'Partitions(replica): -\n' +
                           'TabletTypes: master,rdonly,replica')

ALL_DB_TYPES = ['master', 'replica', 'rdonly']

class TestKeyspace(unittest.TestCase):
  def _read_keyspace(self, keyspace_name):
    global vtgate_port
    vtgate_client = zkocc.ZkOccConnection("localhost:%u" % vtgate_port,
                                        "test_nj", 30.0)
    return keyspace.read_keyspace(vtgate_client, keyspace_name)

  def test_shard_count(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    self.assertEqual(sharded_ks.shard_count, 2)
    for db_type in ALL_DB_TYPES:
      self.assertEqual(sharded_ks.get_shard_count(db_type), 2)
    unsharded_ks = self._read_keyspace(UNSHARDED_KEYSPACE)
    self.assertEqual(unsharded_ks.shard_count, 1)
    for db_type in ALL_DB_TYPES:
      self.assertEqual(unsharded_ks.get_shard_count(db_type), 1)

  def test_shard_names(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    self.assertEqual(sharded_ks.shard_names, ['-80', '80-'])
    for db_type in ALL_DB_TYPES:
      self.assertEqual(sharded_ks.get_shard_names(db_type), ['-80', '80-'])
    unsharded_ks = self._read_keyspace(UNSHARDED_KEYSPACE)
    self.assertEqual(unsharded_ks.shard_names, ['0'])
    for db_type in ALL_DB_TYPES:
      self.assertEqual(unsharded_ks.get_shard_names(db_type), ['0'])

  def test_shard_max_keys(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    want = ['80', '']
    for i, smk in enumerate(sharded_ks.shard_max_keys):
      self.assertEqual(smk.encode('hex').upper(), want[i])
    for db_type in ALL_DB_TYPES:
      for i, smk in enumerate(sharded_ks.get_shard_max_keys(db_type)):
        self.assertEqual(smk.encode('hex').upper(), want[i])
    unsharded_ks = self._read_keyspace(UNSHARDED_KEYSPACE)
    self.assertEqual(unsharded_ks.shard_max_keys, None)
    for db_type in ALL_DB_TYPES:
      self.assertEqual(unsharded_ks.get_shard_max_keys(db_type), [''])

  def test_db_types(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    self.assertEqual(set(sharded_ks.db_types), set(ALL_DB_TYPES))
    unsharded_ks = self._read_keyspace(UNSHARDED_KEYSPACE)
    self.assertEqual(set(unsharded_ks.db_types), set(ALL_DB_TYPES))

  def test_keyspace_id_to_shard_index(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    for i, sn in enumerate(shard_names):
      for keyspace_id in shard_kid_map[sn]:
        self.assertEqual(sharded_ks.keyspace_id_to_shard_index(keyspace_id), i)
        self.assertEqual(sharded_ks.keyspace_id_to_shard_index_for_db_type(keyspace_id, 'master'), i)

  def test_keyspace_id_to_shard_name(self):
    sharded_ks = self._read_keyspace(SHARDED_KEYSPACE)
    for _, sn in enumerate(shard_names):
      for keyspace_id in shard_kid_map[sn]:
        self.assertEqual(sharded_ks.keyspace_id_to_shard_name_for_db_type(keyspace_id, 'master'), sn)
    unsharded_ks = self._read_keyspace(UNSHARDED_KEYSPACE)
    for keyspace_id in shard_kid_map[sn]:
      self.assertEqual(unsharded_ks.keyspace_id_to_shard_name_for_db_type(keyspace_id, 'master'), '0')

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


<commit_before># -*- coding: utf-8 -*-
# Copyright (c) 2012 Thomas Parslow http://almostobsolete.net/
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish, dis-
# tribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the fol-
# lowing conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#

import urllib
import hashlib
import math
import json

def chunk_hashes(str):
  """
  Break up the byte-string into 1MB chunks and return sha256 hashes
  for each.
  """
  chunk = 1024*1024
  chunk_count = int(math.ceil(len(str)/float(chunk)))
  chunks = [str[i*chunk:(i+1)*chunk] for i in range(chunk_count)]
  return [hashlib.sha256(x).digest() for x in chunks]

def tree_hash(hashes):
  """
  Given a hash of each 1MB chunk (from chunk_hashes) this will hash
  together adjacent hashes until it ends up with one big one. So a
  tree of hashes.
  """
  while len(hashes) > 1:
      hashes = [hashlib.sha256("".join(h[i:i+1])).digest() for i in range(i,2)]
  return hashes[0]

def bytes_to_hex(str):
  return ''.join( [ "%02x" % ord( x ) for x in str] ).strip()

class Writer(object):
  """
  Presents a file-like object for writing to a Amazon Glacier
  Archive. The data is written using the multi-part upload API.
  """
  def __init__(self, vault, upload_id, part_size):
      self.vault = vault
      self.upload_id = upload_id
      self.part_size = part_size
      self.buffer_size = 0
      self.uploaded_size = 0
      self.buffer = []
      self.vault = vault
      self.tree_hashes = []
      self.archive_location = None
      self.closed = False

  def make_request(self, verb, headers=None,
                   data='', ok_responses=(200,)):
      resource = "multipart-uploads/%s" % (urllib.quote(self.upload_id),)
      return self.vault.make_request(verb, resource, headers, data, ok_responses)

  def send_part(self):
      buf = "".join(self.buffer)
      # Put back any data remaining over the part size into the
      # buffer
      if len(buf) < self.part_size:
          self.buffer = [buf[self.part_size:]]
          self.buffer_size = len(self.buffer[0])
      else:
          self.buffer = []
          self.buffer_size = 0
      # The part we will send
      part = buf[:self.part_size]
      # Create a request and sign it
      part_tree_hash = tree_hash(chunk_hashes(part))
      self.tree_hashes.append(part_tree_hash)

      headers = {
                  "Content-Range": "bytes %d-%d/*" % (self.uploaded_size, (self.uploaded_size+len(part))-1),
                  "Content-Length": str(len(part)),
                  "Content-Type": "application/octet-stream",
                  "x-amz-sha256-tree-hash": bytes_to_hex(part_tree_hash),
                  "x-amz-content-sha256": hashlib.sha256(part).hexdigest()
                }

      repsonse = self.make_request("PUT", headers, part, ok_responses=(204,))

      self.uploaded_size += len(part)

  def write(self, str):
      assert not self.closed, "Tried to write to a Writer that is already closed!"
      self.buffer.append(str)
      self.buffer_size += len(str)
      while self.buffer_size > self.part_size:
          self.send_part()

  def close(self):
      if self.closed:
          return
      if self.buffer_size > 0:
          self.send_part()
      # Complete the multiplart glacier upload
      headers = {
                  "x-amz-sha256-tree-hash": bytes_to_hex(tree_hash(self.tree_hashes)),
                  "x-amz-archive-size": str(self.uploaded_size)
                }
        response = self.make_request("POST", headers, ok_responses=(201,))
        self.archive_id = response.getheader("x-amz-archive-id")
        self.closed = True

    def get_archive_id(self):
        self.close()
        return self.archive_id
<commit_msg>Removed additional indentation that got added just before the last commit<commit_after># -*- coding: utf-8 -*-
# Copyright (c) 2012 Thomas Parslow http://almostobsolete.net/
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish, dis-
# tribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the fol-
# lowing conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#

import urllib
import hashlib
import math
import json

def chunk_hashes(str):
  """
  Break up the byte-string into 1MB chunks and return sha256 hashes
  for each.
  """
  chunk = 1024*1024
  chunk_count = int(math.ceil(len(str)/float(chunk)))
  chunks = [str[i*chunk:(i+1)*chunk] for i in range(chunk_count)]
  return [hashlib.sha256(x).digest() for x in chunks]

def tree_hash(hashes):
  """
  Given a hash of each 1MB chunk (from chunk_hashes) this will hash
  together adjacent hashes until it ends up with one big one. So a
  tree of hashes.
  """
  while len(hashes) > 1:
      hashes = [hashlib.sha256("".join(h[i:i+1])).digest() for i in range(i,2)]
  return hashes[0]

def bytes_to_hex(str):
  return ''.join( [ "%02x" % ord( x ) for x in str] ).strip()

class Writer(object):
  """
  Presents a file-like object for writing to a Amazon Glacier
  Archive. The data is written using the multi-part upload API.
  """
  def __init__(self, vault, upload_id, part_size):
      self.vault = vault
      self.upload_id = upload_id
      self.part_size = part_size
      self.buffer_size = 0
      self.uploaded_size = 0
      self.buffer = []
      self.vault = vault
      self.tree_hashes = []
      self.archive_location = None
      self.closed = False

  def make_request(self, verb, headers=None,
                   data='', ok_responses=(200,)):
      resource = "multipart-uploads/%s" % (urllib.quote(self.upload_id),)
      return self.vault.make_request(verb, resource, headers, data, ok_responses)

  def send_part(self):
      buf = "".join(self.buffer)
      # Put back any data remaining over the part size into the
      # buffer
      if len(buf) < self.part_size:
          self.buffer = [buf[self.part_size:]]
          self.buffer_size = len(self.buffer[0])
      else:
          self.buffer = []
          self.buffer_size = 0
      # The part we will send
      part = buf[:self.part_size]
      # Create a request and sign it
      part_tree_hash = tree_hash(chunk_hashes(part))
      self.tree_hashes.append(part_tree_hash)

      headers = {
                  "Content-Range": "bytes %d-%d/*" % (self.uploaded_size, (self.uploaded_size+len(part))-1),
                  "Content-Length": str(len(part)),
                  "Content-Type": "application/octet-stream",
                  "x-amz-sha256-tree-hash": bytes_to_hex(part_tree_hash),
                  "x-amz-content-sha256": hashlib.sha256(part).hexdigest()
                }

      repsonse = self.make_request("PUT", headers, part, ok_responses=(204,))

      self.uploaded_size += len(part)

  def write(self, str):
      assert not self.closed, "Tried to write to a Writer that is already closed!"
      self.buffer.append(str)
      self.buffer_size += len(str)
      while self.buffer_size > self.part_size:
          self.send_part()

  def close(self):
      if self.closed:
          return
      if self.buffer_size > 0:
          self.send_part()
      # Complete the multiplart glacier upload
      headers = {
                  "x-amz-sha256-tree-hash": bytes_to_hex(tree_hash(self.tree_hashes)),
                  "x-amz-archive-size": str(self.uploaded_size)
                }
      response = self.make_request("POST", headers, ok_responses=(201,))
      self.archive_id = response.getheader("x-amz-archive-id")
      self.closed = True

  def get_archive_id(self):
      self.close()
      return self.archive_id
<|endoftext|>


<commit_before>import yaml

def load_config(config_path):
    <commit_msg>first commit<commit_after>from __future__ import print_function

import sys
import time
import yaml
import os
import os.path
import logging
import shutil

LOG = logging.getLogger(__name__)

COMMITTED_CONFIG = 'committed-config.yml'
RUNNING_CONFIG = 'running-config.yml'
PROTOTYPE_ENV = 'MINEMELD_PROTOTYPE_PATH'

def _load_node_prototype(protoname):
    paths = os.getenv(PROTOTYPE_ENV, None)
    if paths is None:
        raise RuntimeError('Unable to load prototype %s: %s '
                           'environment variable not set' %
                           (protoname, PROTOTYPE_ENV))
    paths = paths.split(':')

    proto_module, proto_name = protoname.rsplit('.', 1)

    pmodule = None
    pmprotos = {}
    for p in paths:
        pmpath = os.path.join(p, proto_module+'.yml')

        try:
            with open(pmpath, 'r') as pf:
                pmodule = yaml.safe_load(pf)

                if pmodule is None:
                    pmodule = {}
        except IOError:
            pmodule = None
            continue

        pmprotos = pmodule.get('prototypes', {})

        if proto_name not in pmprotos:
            pmodule = None
            continue

        if 'class' not in pmprotos[proto_name]:
            pmodule = None
            continue

        return pmprotos[proto_name]

    raise RuntimeError('Unable to load prototype %s: '
                       ' not found' % (protoname))

def _load_config_from_file(f):
    config = yaml.safe_load(f)

    if config is None:
        config = {}

    if 'fabric' not in config:
        config['fabric'] = {
            'class': 'AMQP',
            'config': {
                'num_connections': 5
            }
        }

    if 'mgmtbus' not in config:
        config['mgmtbus'] = {
            'transport': {
                'class': 'AMQP',
                'config': {}
            },
            'master': {},
            'slave': {}
        }

    nodes_config = config.get('nodes', {})
    for nname, nconfig in nodes_config.iteritems():
        if 'prototype' in nconfig:
            nproto = _load_node_prototype(nconfig['prototype'])

            nconfig.pop('prototype')

            nconfig['class'] = nproto['class']
            nproto_config = nproto.get('config', {})
            nproto_config.update(
                nconfig.get('config', {})
            )
            nconfig['config'] = nproto_config

    return config

def _load_config_from_dir(path):
    ccpath = os.path.join(
        path,
        COMMITTED_CONFIG
    )
    rcpath = os.path.join(
        path,
        RUNNING_CONFIG
    )

    cconfig = None
    if os.path.exists(ccpath):
        with open(ccpath, 'r') as cf:
            cconfig = _load_config_from_file(cf)

    rcconfig = None
    if os.path.exists(rcpath):
        with open(rcpath, 'r') as cf:
            rcconfig = _load_config_from_file(cf)

    if rcconfig is None and cconfig is None:
        print(
            "At least one of", RUNNING_CONFIG,
            "or", COMMITTED_CONFIG,
            "should exist in", path,
            file=sys.stderr
        )
        sys.exit(1)
    elif rcconfig is not None and cconfig is None:
        rcconfig['newconfig'] = False
        return rcconfig
    elif rcconfig is None and cconfig is not None:
        shutil.copyfile(ccpath, rcpath)
        cconfig['newconfig'] = True
        return cconfig
    elif rcconfig is not None and cconfig is not None:
        # ugly
        if yaml.dump(cconfig) != yaml.dump(rcconfig):
            shutil.copyfile(rcpath, rcpath+'.%d' % int(time.time()))
            shutil.copyfile(ccpath, rcpath)
            cconfig['newconfig'] = True
            return cconfig

        rcconfig['newconfig'] = False
        return rcconfig

def load_config(config_path):
    if os.path.isdir(config_path):
        return _load_config_from_dir(config_path)

    with open(config_path, 'r') as cf:
        config = _load_config_from_file(cf)

    config['newconfig'] = True

    return config
<|endoftext|>


<commit_before>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Search and Store Tweets utility.

Search for tweets in the Twitter API for given input terms, then store
the tweet and the tweet's author data locally, updating or adding objects
as required. The input may be command-line text for easy ad hoc queries,
or the name of an existing Campaign in the db so that its search query string
can be looked up.

Add the default label as Campaign for processed Tweets to signal that the Tweet
was processed by this script. If looking up Tweets using the search query from
a Campaign specified in arguments, then also allocate the processed Tweets
to that Campaign.

Add the default label as Category for Profiles which are processed due to
processing their Tweets.

The persist value is set based on an answer here:
    https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse

TODO: Consolidate use of writeToDB and persist in this repo.
"""
import argparse
import datetime
import os
import sys

from sqlobject import SQLObjectNotFound

# Allow imports to be done when executing this file directly.
appDir = os.path.abspath(os.path.join(os.path.dirname(__file__),
                                      os.path.pardir, os.path.pardir))
sys.path.insert(0, appDir)

from lib import database as db, flattenText, tweets
from lib.twitter import auth, search
from lib.query.tweets.campaigns import printAvailableCampaigns,\
                                       printCampaignsAndTweets

# Name to assign to Tweets and Profiles created or updated with this utility.
BASE_LABEL = u"_SEARCH_QUERY"
# Setup global API connection object, which needs to be set using a function
# on auth.
API_CONN = None

def searchAndStore(searchQuery, totalCount=200, persist=True,
                   extended=True):
    """
    Search the Twitter Search API for tweets matching input search terms.

    By default, the tweets are created or updated as Tweet records in the local
    db.

    Only matches on tweets for users which had their language set to English
    or undefined.

    @param searchQuery: query text to search on the Twitter API.
    @param totalCount: total count of tweets to attempt to get for the
        search query, as an integer. Defaults to 200, which is the max count
        of tweets received on a single page from the Twitter API.
    @param persist. Default True. If set to False, does not store data
        in the database and only prints to stdout.
    @param extended: If True, get the expanded tweet message instead of the
        truncated form.

    @return processedTweets: count of tweets fetched, unaffected by
        with the data is persisted. This count will be a number up to the
        totalCount argument, but may less if fewer tweets are available in
        the 7-day window.
    @return profileRecs: List of local Profile records inserted or updated.
        Defaults to empty list.
    @return tweetRecs: List of local Tweet records inserted or updated.
        Defaults to empty list.
    """
    assert API_CONN, ("Authenticate with Twitter API before doing"
                      " a search for tweets.")
    searchResults = search.fetchTweetsPaging(
        API_CONN,
        searchQuery=searchQuery,
        itemLimit=totalCount,
        extended=extended
    )

    processedTweets = 0
    profileRecs = []
    tweetRecs = []

    for fetchedTweet in searchResults:
        if persist:
            profileRec = tweets.insertOrUpdateProfile(fetchedTweet.author)
            profileRecs.append(profileRec)
            data, tweetRec = tweets.insertOrUpdateTweet(
                fetchedTweet,
                profileRec.id
            )
            tweetRecs.append(tweetRec)
            if (processedTweets + 1) % 100 == 0:
                print "Processed so far: {}".format(processedTweets + 1)
        else:
            # Assume extended mode, otherwise fall back to standard mode.
            try:
                text = fetchedTweet.full_text
            except AttributeError:
                text = fetchedTweet.text

            print u"{index:3d} @{screenName}: {message}".format(
                index=processedTweets + 1,
                screenName=fetchedTweet.author.screen_name,
                message=flattenText(text)
            )
        processedTweets += 1
    print

    return processedTweets, profileRecs, tweetRecs

def main():
    """
    Handle command-line arguments to search for tweets, store data for
    Tweet and Profile objects and then assign Campaigns to Tweets.
    """
    global API_CONN

    parser = argparse.ArgumentParser(
        description="""\
Utility to search for tweets and then the store tweet and profile data locally.

Search with either an ad hoc query, or the name of a stored Campaign which
has a search query set. To create or edit a Campaign, use the Campaign Manager
utility.
""",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    view = parser.add_argument_group("View", "Print data to stdout")
    view.add_argument(
        '-a', '--available',
        action='store_true',
        help="""Output available Campaigns in db, with Tweet counts and
             search query for each."""
    )
    view.add_argument(
        '-t', '--tweets',
        action='store_true',
        help="Output local Tweets grouped by Campaign."
    )
    view.add_argument(
        '-s', '--search-help',
        action='store_true',
        help="""Print guide for writing search queries, with examples of
            syntax safe for the command-line. See Twitter's search
            documentation for full rules."""
    )

    fetch = parser.add_argument_group("Fetch", "Select a search query to"
                                      " get Tweets from Twitter Search API.")
    fetch.add_argument(
        '-c', '--campaign',
        help="""Name of existing campaign in the db. If supplied and the
            Campaign record has a query string, fetch Tweets from the Twitter
            Search API and store, assigning the Campaign name and
            {baseLabel} as campaigns on all processed Tweets.
            This argument may not be used with the --query argument.
        """.format(baseLabel=BASE_LABEL)
    )
    fetch.add_argument(
        '-q', '--query',
        help="""Word or phrase to search on the Twitter API as an ad hoc
            query which is not associated with a Campaign. This argument
            may not be used with the --campaign argument."""
    )
    fetch.add_argument(
        '-C', '--count',
        metavar='N',
        type=int,
        default=200,
        help="Default 200. Max count of tweets to get for the search query.")
    fetch.add_argument(
        '--persist',
        dest='persist',
        action='store_true',
        help="(DEFAULT) Store fetched tweets and profiles in the database.")
    fetch.add_argument(
        '--no-persist',
        dest='persist',
        action='store_false',
        help="Print fetched tweet and profile datas without storing.")
    fetch.set_defaults(persist=True)

    args = parser.parse_args()

    if args.available:
        printAvailableCampaigns()
    if args.tweets:
        printCampaignsAndTweets()
    if args.search_help:
        print search.getSearchQueryHelp()

    if args.query or args.campaign:
        try:
            generalCampaignRec = db.Campaign.byName(BASE_LABEL)
        except SQLObjectNotFound:
            # The campaign manager is not needed externally for creating
            # this one, since the searchQuery is best set to NULL for
            # this specific campaign and therefore can be automatic.
            generalCampaignRec = db.Campaign(name=BASE_LABEL, searchQuery=None)

        if args.query:
            customCampaignRec = None
            query = unicode(args.query)
        else:
            campaignName = args.campaign
            try:
                customCampaignRec = db.Campaign.byName(campaignName)
            except SQLObjectNotFound as e:
                raise type(e)("Use the campaign manager to create the Campaign"
                              " as name and search query. Name not found: {0}"
                              .format(campaignName))
            query = customCampaignRec.searchQuery
            assert query, "Use the Campaign Mananger to set a search query"\
                          " for the campaign: {0}".format(args.campaign)

        print u"Search query: {0}".format(query)

        # Use app auth  herefor up to 450 search requests per window, rather
        # than 180.
        API_CONN = auth.getAppOnlyConnection()

        now = datetime.datetime.now()
        processedCount, profileRecs, tweetRecs = searchAndStore(
            query,
            totalCount=args.count,
            persist=args.persist
        )
        print "Completed tweet processing: {0:,d}".format(processedCount)
        print "took {}".format(datetime.datetime.now()-now)

        if profileRecs:
            print "Assigning category links... ",
            now = datetime.datetime.now()
            # TODO: Setup as default data in database.py rather than
            # create here.
            try:
                generalCategoryRec = db.Category.byName(BASE_LABEL)
            except SQLObjectNotFound:
                generalCategoryRec = db.Category(name=BASE_LABEL)
            tweets.bulkAssignProfileCategory(
                categoryID=generalCategoryRec.id,
                profileIDs=(profile.id for profile in profileRecs)
            )
            print "DONE"
            print "took {}".format(datetime.datetime.now()-now)

        if tweetRecs:
            print "Assigning general campaign links... ",
            now = datetime.datetime.now()
            tweets.bulkAssignTweetCampaign(
                campaignID=generalCampaignRec.id,
                tweetIDs=(tweet.id for tweet in tweetRecs)
            )
            print "DONE"
            print "took {}".format(datetime.datetime.now()-now)

            if customCampaignRec:
                print "Assigning custom campaign links... ",
                now = datetime.datetime.now()
                # Reset generator to first item, after using it above within
                # the bulk assign function.
                tweetIDs = (tweet.id for tweet in tweetRecs)
                tweets.bulkAssignTweetCampaign(
                    campaignID=customCampaignRec.id,
                    tweetIDs=tweetIDs
                )
                print "DONE"
                print "took {}".format(datetime.datetime.now()-now)

if __name__ == '__main__':
    main()
<commit_msg>feature: Add configured labels to searchAndStoreTweets.py style: Update searchAndStoreTweets.py<commit_after>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Search and Store Tweets utility.

Search for tweets in the Twitter API for given input terms, then store
the tweet and the tweet's author data locally, updating or adding objects
as required. The input may be command-line text for easy ad hoc queries,
or the name of an existing Campaign in the db so that its search query string
can be looked up.

Add the default label as Campaign for processed Tweets to signal that the Tweet
was processed by this script. If looking up Tweets using the search query from
a Campaign specified in arguments, then also allocate the processed Tweets
to that Campaign.

Add the default label as Category for Profiles which are processed due to
processing their Tweets.

The persist value is set based on an answer here:
    https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse

TODO: Consolidate use of writeToDB and persist in this repo.
"""
import argparse
import datetime
import os
import sys

from sqlobject import SQLObjectNotFound

# Allow imports to be done when executing this file directly.
sys.path.insert(0, os.path.abspath(os.path.join(
    os.path.dirname(__file__), os.path.pardir, os.path.pardir)
))

from lib import database as db, flattenText, tweets
from lib.config import AppConf
from lib.twitter import auth, search
from lib.query.tweets.campaigns import printAvailableCampaigns, \
                                       printCampaignsAndTweets

conf = AppConf ()
UTILITY_CATEGORY = UTILITY_CAMPAIGN = conf.get('Labels', 'search')

# Create initial global API connection object, which needs to be set using
# a function on auth.
API_CONN = None

def searchAndStore(searchQuery, totalCount=200, persist=True,
                   extended=True):
    """
    Search the Twitter Search API for tweets matching input search terms.

    By default, the tweets are created or updated as Tweet records in the local
    db.

    Only matches on tweets for users which had their language set to English
    or undefined.

    @param searchQuery: query text to search on the Twitter API.
    @param totalCount: total count of tweets to attempt to get for the
        search query, as an integer. Defaults to 200, which is the max count
        of tweets received on a single page from the Twitter API.
    @param persist. Default True. If set to False, does not store data
        in the database and only prints to stdout.
    @param extended: If True, get the expanded tweet message instead of the
        truncated form.

    @return processedTweets: count of tweets fetched, unaffected by
        with the data is persisted. This count will be a number up to the
        totalCount argument, but may less if fewer tweets are available in
        the 7-day window.
    @return profileRecs: List of local Profile records inserted or updated.
        Defaults to empty list.
    @return tweetRecs: List of local Tweet records inserted or updated.
        Defaults to empty list.
    """
    assert API_CONN, ("Authenticate with Twitter API before doing"
                      " a search for tweets.")
    searchResults = search.fetchTweetsPaging(
        API_CONN,
        searchQuery=searchQuery,
        itemLimit=totalCount,
        extended=extended
    )

    processedTweets = 0
    profileRecs = []
    tweetRecs = []

    for fetchedTweet in searchResults:
        if persist:
            profileRec = tweets.insertOrUpdateProfile(fetchedTweet.author)
            profileRecs.append(profileRec)
            data, tweetRec = tweets.insertOrUpdateTweet(
                fetchedTweet,
                profileRec.id
            )
            tweetRecs.append(tweetRec)
            if (processedTweets + 1) % 100 == 0:
                print "Processed so far: {}".format(processedTweets + 1)
        else:
            # Assume extended mode, otherwise fall back to standard mode.
            try:
                text = fetchedTweet.full_text
            except AttributeError:
                text = fetchedTweet.text

            print u"{index:3d} @{screenName}: {message}".format(
                index=processedTweets + 1,
                screenName=fetchedTweet.author.screen_name,
                message=flattenText(text)
            )
        processedTweets += 1
    print

    return processedTweets, profileRecs, tweetRecs

def main():
    """
    Handle command-line arguments to search for tweets, store data for
    Tweet and Profile objects and then assign Campaigns to Tweets.
    """
    global API_CONN

    parser = argparse.ArgumentParser(
        description="""\
Utility to search for tweets and then the store tweet and profile data locally.

Search with either an ad hoc query, or the name of a stored Campaign which
has a search query set. To create or edit a Campaign, use the Campaign Manager
utility.
""",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    view = parser.add_argument_group(
        "View",
        "Print data to stdout"
    )
    view.add_argument(
        '-a', '--available',
        action='store_true',
        help="""Output available Campaigns in db, with Tweet counts and
             search query for each."""
    )
    view.add_argument(
        '-t', '--tweets',
        action='store_true',
        help="Output local Tweets grouped by Campaign."
    )
    view.add_argument(
        '-s', '--search-help',
        action='store_true',
        help="""Print guide for writing search queries, with examples of
            syntax safe for the command-line. See Twitter's search
            documentation for full rules."""
    )

    fetch = parser.add_argument_group(
        "Fetch",
        "Select a search query to get Tweets from Twitter Search API."
    )
    fetch.add_argument(
        '-c', '--campaign',
        help="""Name of existing campaign in the db. If supplied and the
            Campaign record has a query string, fetch Tweets from the Twitter
            Search API and store. Then assign the given custom Campaign name
            to processed Tweets. This argument may not be used with the
            --query argument.
        """
    )
    fetch.add_argument(
        '-q', '--query',
        help="""Word or phrase to search on the Twitter API as an ad hoc
            query which is not associated with a Campaign. This argument
            may not be used with the --campaign argument."""
    )
    fetch.add_argument(
        '-C', '--count',
        metavar='N',
        type=int,
        default=200,
        help="Default 200. Max count of tweets to get for the search query.")
    fetch.add_argument(
        '--persist',
        dest='persist',
        action='store_true',
        help="(DEFAULT) Store fetched tweets and profiles in the database.")
    fetch.add_argument(
        '--no-persist',
        dest='persist',
        action='store_false',
        help="Print fetched tweet and profile datas without storing.")
    fetch.set_defaults(persist=True)

    args = parser.parse_args()

    if args.available:
        printAvailableCampaigns()
    if args.tweets:
        printCampaignsAndTweets()
    if args.search_help:
        print search.getSearchQueryHelp()

    if args.query or args.campaign:
        try:
            utilityCampaignRec = db.Campaign.byName(UTILITY_CAMPAIGN)
        except SQLObjectNotFound:
            # The campaign manager is not needed externally for creating
            # this one, since the searchQuery is best set to NULL for
            # this specific campaign and therefore can be automatic.
            utilityCampaignRec = db.Campaign(
                name=UTILITY_CAMPAIGN,
                searchQuery=None
            )

        if args.query:
            customCampaignRec = None
            query = unicode(args.query, 'utf-8')
        else:
            campaignName = args.campaign
            try:
                customCampaignRec = db.Campaign.byName(campaignName)
            except SQLObjectNotFound as e:
                raise type(e)("Use the campaign manager to create the Campaign"
                              " as name and search query. Name not found: {0}"
                              .format(campaignName))
            query = customCampaignRec.searchQuery
            assert query, "Use the Campaign Mananger to set a search query"\
                          " for the campaign: {0}".format(args.campaign)

        # Process the category and campaign records above before fetching
        # data from the API.
        print u"Search query: {0}".format(query)

        # Use app auth here for up to 480 search requests per window, rather
        # than 180 when using the user auth.
        API_CONN = auth.getAppOnlyConnection()

        now = datetime.datetime.now()
        processedCount, profileRecs, tweetRecs = searchAndStore(
            query,
            totalCount=args.count,
            persist=args.persist
        )
        print "Completed tweet processing: {0:,d}".format(processedCount)
        print "took {0}".format(datetime.datetime.now() - now)

        if profileRecs:
            print "Assigning category links... ",
            now = datetime.datetime.now()
            try:
                utilityCategoryRec = db.Category.byName(UTILITY_CATEGORY)
            except SQLObjectNotFound:
                utilityCategoryRec = db.Category(name=UTILITY_CATEGORY)
            tweets.bulkAssignProfileCategory(
                categoryID=utilityCategoryRec.id,
                profileIDs=(profile.id for profile in profileRecs)
            )
            print "DONE"
            print "took {0}".format(datetime.datetime.now() - now)

        if tweetRecs:
            print "Assigning utility's campaign links... ",
            now = datetime.datetime.now()
            tweets.bulkAssignTweetCampaign(
                campaignID=utilityCampaignRec.id,
                tweetIDs=(tweet.id for tweet in tweetRecs)
            )
            print "DONE"
            print "took {0}".format(datetime.datetime.now() - now)

            if customCampaignRec:
                print "Assigning custom campaign links... ",
                now = datetime.datetime.now()
                # Reset generator to first item, after using it above within
                # the bulk assign function.
                tweetIDs = (tweet.id for tweet in tweetRecs)
                tweets.bulkAssignTweetCampaign(
                    campaignID=customCampaignRec.id,
                    tweetIDs=tweetIDs
                )
                print "DONE"
                print "took {0}".format(datetime.datetime.now() - now)

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


<commit_before>from django.views.generic import DetailView, ListView
from django.views.generic import DayArchiveView, MonthArchiveView
from django.views.generic import YearArchiveView

from .models import Entry

class TaggedEntryListView(ListView):
    """List all ``Entry``s that have the given ``Tag``(s). Mulitple ``Tag``s
    may be separated by a plus; For example: /blog/tags/foo+bar would retrieve
    all ``Entry``s tagged with both "foo" and "bar".

    """
    allow_empty = True
    model = Entry

    def get_queryset(self):
        tag_list = self.kwargs['tag_slug'].split('+')
        tags = filter(lambda x: len(x) > 0, tag_list)
        return Entry.objects.filter(tags__slug__in=tags).distinct()

class EntryDetailView(DetailView):
    """Detail for an ``Entry``."""
    model = Entry
    slug_field = 'slug'

# Year, Month, Day Archives
# -------------------------

class EntryYearArchiveView(YearArchiveView):
    queryset = Entry.objects.filter(published=True)
    date_field = "published_on"
    year_format = '%Y'
    template_name = "blargg/entry_archive_year.html"

class EntryMonthArchiveView(MonthArchiveView):
    queryset = Entry.objects.filter(published=True)
    date_field = "published_on"
    year_format = '%Y'
    month_format = "%m"
    template_name = "blargg/entry_archive_month.html"

class EntryDayArchiveView(DayArchiveView):
    # NOTE: Entries are stored in UTC and this view converts dates to the
    # local timezone (if USE_TZ=True). Therefore, Entry.get_absolute_url also
    # converts to TIME_ZONE if USE_TZ=True.
    queryset = Entry.objects.filter(published=True)
    date_field = "published_on"
    year_format = '%Y'
    month_format = "%m"
    day_format = "%d"
    template_name = "blargg/entry_archive_day.html"
    allow_empty = True
<commit_msg>added an EntryStatsMixin that adds entry stats to the context of entries<commit_after>from django.views.generic import DetailView, ListView
from django.views.generic import DayArchiveView, MonthArchiveView
from django.views.generic import YearArchiveView
from django.views.generic.list import MultipleObjectMixin

from .models import Entry, entry_stats

class EntryStatsMixin(MultipleObjectMixin):
    def get_context_data(self, **kwargs):
        context = super(EntryStatsMixin, self).get_context_data(**kwargs)
        # Calculate stats for the queryset of Entries & add to the context.
        context.update(entry_stats(context['object_list']))
        return context

class TaggedEntryListView(ListView):
    """List all ``Entry``s that have the given ``Tag``(s). Mulitple ``Tag``s
    may be separated by a plus; For example: /blog/tags/foo+bar would retrieve
    all ``Entry``s tagged with both "foo" and "bar".

    """
    allow_empty = True
    model = Entry

    def get_queryset(self):
        tag_list = self.kwargs['tag_slug'].split('+')
        tags = filter(lambda x: len(x) > 0, tag_list)
        return Entry.objects.filter(tags__slug__in=tags).distinct()

class EntryDetailView(DetailView):
    """Detail for an ``Entry``."""
    model = Entry
    slug_field = 'slug'

# Year, Month, Day Archives
# -------------------------

class EntryYearArchiveView(EntryStatsMixin, YearArchiveView):
    queryset = Entry.objects.filter(published=True)
    date_field = "published_on"
    year_format = '%Y'
    template_name = "blargg/entry_archive_year.html"

class EntryMonthArchiveView(EntryStatsMixin, MonthArchiveView):
    queryset = Entry.objects.filter(published=True)
    date_field = "published_on"
    year_format = '%Y'
    month_format = "%m"
    template_name = "blargg/entry_archive_month.html"

class EntryDayArchiveView(EntryStatsMixin, DayArchiveView):
    # NOTE: Entries are stored in UTC and this view converts dates to the
    # local timezone (if USE_TZ=True). Therefore, Entry.get_absolute_url also
    # converts to TIME_ZONE if USE_TZ=True.
    queryset = Entry.objects.filter(published=True)
    date_field = "published_on"
    year_format = '%Y'
    month_format = "%m"
    day_format = "%d"
    template_name = "blargg/entry_archive_day.html"
    allow_empty = True
<|endoftext|>