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

import argparse
import json
import logging
import os
import shutil
import time
import re
import sys
import select
import termios

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

class Color:
    BLACK = '\033[90m'
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    PINK = '\033[95m'
    CYAN = '\033[96m'
    WHITE = '\033[97m'
    ENDC = '\033[0m'

COLOR_FOUND = Color.Pink
COLOR_MONITOR = Color.CYAN

class Status:
    NEW_TORRENTFILE_FOUND = 0
    MONITOR = 1

status_messages = {
    Status.NEW_TORRENTFILE_FOUND: '%FOUND%s' % (COLOR_FOUND, Color.ENDC),
    Status.MONITOR: '%MONITOR%s' % (COLOR_MONITOR, Color.ENDC),
}

class KeyPoller():
    def __enter__(self):
        # Save the terminal settings
        self.fd = sys.stdin.fileno()
        self.new_term = termios.tcgetattr(self.fd)
        self.old_term = termios.tcgetattr(self.fd)

        # New terminal setting unbuffered
        self.new_term[3] = (self.new_term[3] & ~termios.ICANON & ~termios.ECHO)
        termios.tcsetattr(self.fd, termios.TCSAFLUSH, self.new_term)

        return self

    def __exit__(self, type, value, traceback):
        termios.tcsetattr(self.fd, termios.TCSAFLUSH, self.old_term)

    def poll(self):
        dr,dw,de = select.select([sys.stdin], [], [], 0)
        if not dr == []:
            return sys.stdin.read(1)
        return None

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():
    print('###### autotorrent-1.6.2e1 ######')
    print('# Original by John Doee https://github.com/JohnDoee/autotorrent')
    print('# Monitoring mode added by Jean-Francois Drapeau https://github.com/jeanfrancoisdrapeau/autotorrent')

    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_status(Status.MONITOR, args.loopmode, '(press X to exit)')
        with KeyPoller() as keyPoller:
            while True:
                c = keyPoller.poll()
                if not c is None:
                    if c == "x":
                        break

                for fn in os.listdir(args.loopmode):
                    if fn.endswith('.torrent'):
                        db.rebuild([config.get('general', 'store_path')])

                        fn_woext = os.path.splitext(fn)[0]
                        fn_scenename = re.search('-(.*)$', fn_woext).group(1).replace(' ', '.').lower()
                        print_status(Status.NEW_TORRENTFILE_FOUND, fn_woext, 'New torrent added successfully')

                        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, False)

                                    # 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")
                            addtfile(at, args.loopmode, [fn], args.dry_run, True)

                            # 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')])

                        print_status(Status.MONITOR, args.loopmode, '(press X to exit)')

                time.sleep(5)

    print('Goodbye!')

def addtfile(at, current_path, afiles, adry_run, is_new):
    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, is_new)
        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('')

def print_status(status, torrentfile, message):
    print(' %-20s %r %s' % ('[%s]' % status_messages[status], os.path.splitext(os.path.basename(torrentfile))[0], message))

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
import sys
import select
import termios

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

class Color:
    BLACK = '\033[90m'
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    PINK = '\033[95m'
    CYAN = '\033[96m'
    WHITE = '\033[97m'
    ENDC = '\033[0m'

COLOR_FOUND = Color.PINK
COLOR_MONITOR = Color.CYAN

class Status:
    NEW_TORRENTFILE_FOUND = 0
    MONITOR = 1

status_messages = {
    Status.NEW_TORRENTFILE_FOUND: '%FOUND%s' % (COLOR_FOUND, Color.ENDC),
    Status.MONITOR: '%MONITOR%s' % (COLOR_MONITOR, Color.ENDC),
}

class KeyPoller():
    def __enter__(self):
        # Save the terminal settings
        self.fd = sys.stdin.fileno()
        self.new_term = termios.tcgetattr(self.fd)
        self.old_term = termios.tcgetattr(self.fd)

        # New terminal setting unbuffered
        self.new_term[3] = (self.new_term[3] & ~termios.ICANON & ~termios.ECHO)
        termios.tcsetattr(self.fd, termios.TCSAFLUSH, self.new_term)

        return self

    def __exit__(self, type, value, traceback):
        termios.tcsetattr(self.fd, termios.TCSAFLUSH, self.old_term)

    def poll(self):
        dr,dw,de = select.select([sys.stdin], [], [], 0)
        if not dr == []:
            return sys.stdin.read(1)
        return None

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():
    print('###### autotorrent-1.6.2e1 ######')
    print('# Original by John Doee https://github.com/JohnDoee/autotorrent')
    print('# Monitoring mode added by Jean-Francois Drapeau https://github.com/jeanfrancoisdrapeau/autotorrent')

    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_status(Status.MONITOR, args.loopmode, '(press X to exit)')
        with KeyPoller() as keyPoller:
            while True:
                c = keyPoller.poll()
                if not c is None:
                    if c == "x":
                        break

                for fn in os.listdir(args.loopmode):
                    if fn.endswith('.torrent'):
                        db.rebuild([config.get('general', 'store_path')])

                        fn_woext = os.path.splitext(fn)[0]
                        fn_scenename = re.search('-(.*)$', fn_woext).group(1).replace(' ', '.').lower()
                        print_status(Status.NEW_TORRENTFILE_FOUND, fn_woext, 'New torrent added successfully')

                        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, False)

                                    # 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")
                            addtfile(at, args.loopmode, [fn], args.dry_run, True)

                            # 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')])

                        print_status(Status.MONITOR, args.loopmode, '(press X to exit)')

                time.sleep(5)

    print('Goodbye!')

def addtfile(at, current_path, afiles, adry_run, is_new):
    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, is_new)
        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('')

def print_status(status, torrentfile, message):
    print(' %-20s %r %s' % ('[%s]' % status_messages[status], os.path.splitext(os.path.basename(torrentfile))[0], message))

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


<commit_before>import warnings

w_s = "The module 'dipy.reconst.peaks' is deprecated."
w_s += " Please use the module 'dipy.direction.peaks' instead"
warnings.warn(w_s, DeprecationWarning)

from dipy.direction.peaks import *
<commit_msg>delete dipy.reconst.peaks module<commit_after><|endoftext|>


<commit_before>from __future__ import absolute_import

from django import forms

from allauth.account.forms import BaseSignupForm
from allauth.account.utils import (user_username, user_email)

from .models import SocialAccount
from .adapter import get_adapter
from . import app_settings
from . import signals

class SignupForm(BaseSignupForm):

    def __init__(self, *args, **kwargs):
        self.sociallogin = kwargs.pop('sociallogin')
        user = self.sociallogin.account.user
        initial = { 'email': user_email(user) or '',
                    'username': user_username(user) or '',
                    'first_name': user.first_name or '',
                    'last_name': user.last_name or '' }
        kwargs['initial'] = initial
        kwargs['email_required'] = app_settings.EMAIL_REQUIRED
        super(SignupForm, self).__init__(*args, **kwargs)

    def save(self, request):
        new_user = self.create_user()
        self.sociallogin.account.user = new_user
        self.sociallogin.save(request)
        super(SignupForm, self).save(new_user)
        return new_user

class DisconnectForm(forms.Form):
    account = forms.ModelChoiceField(queryset=SocialAccount.objects.none(),
                                     widget=forms.RadioSelect,
                                     required=True)

    def __init__(self, *args, **kwargs):
        self.request = kwargs.pop('request')
        self.accounts = SocialAccount.objects.filter(user=self.request.user)
        super(DisconnectForm, self).__init__(*args, **kwargs)
        self.fields['account'].queryset = self.accounts

    def clean(self):
        cleaned_data = super(DisconnectForm, self).clean()
        account = cleaned_data.get('account')
        if account:
            get_adapter().validate_disconnect(account, self.accounts)
        return cleaned_data

    def save(self):
        account = self.cleaned_data['account']
        account.delete()
        signals.social_account_removed.send(sender=SocialAccount,
                                            request=self.request,
                                            socialaccount=account)
<commit_msg>Custom model without first/last name (closes #329)<commit_after>from __future__ import absolute_import

from django import forms

from allauth.account.forms import BaseSignupForm
from allauth.account.utils import (user_username, user_email,
                                   user_field)

from .models import SocialAccount
from .adapter import get_adapter
from . import app_settings
from . import signals

class SignupForm(BaseSignupForm):

    def __init__(self, *args, **kwargs):
        self.sociallogin = kwargs.pop('sociallogin')
        user = self.sociallogin.account.user
        initial = {'email': user_email(user) or '',
                   'username': user_username(user) or '',
                   'first_name': user_field(user, 'first_name') or '',
                   'last_name': user_field(user, 'last_name') or ''}
        kwargs['initial'] = initial
        kwargs['email_required'] = app_settings.EMAIL_REQUIRED
        super(SignupForm, self).__init__(*args, **kwargs)

    def save(self, request):
        new_user = self.create_user()
        self.sociallogin.account.user = new_user
        self.sociallogin.save(request)
        super(SignupForm, self).save(new_user)
        return new_user

class DisconnectForm(forms.Form):
    account = forms.ModelChoiceField(queryset=SocialAccount.objects.none(),
                                     widget=forms.RadioSelect,
                                     required=True)

    def __init__(self, *args, **kwargs):
        self.request = kwargs.pop('request')
        self.accounts = SocialAccount.objects.filter(user=self.request.user)
        super(DisconnectForm, self).__init__(*args, **kwargs)
        self.fields['account'].queryset = self.accounts

    def clean(self):
        cleaned_data = super(DisconnectForm, self).clean()
        account = cleaned_data.get('account')
        if account:
            get_adapter().validate_disconnect(account, self.accounts)
        return cleaned_data

    def save(self):
        account = self.cleaned_data['account']
        account.delete()
        signals.social_account_removed.send(sender=SocialAccount,
                                            request=self.request,
                                            socialaccount=account)
<|endoftext|>


<commit_before><commit_msg>Flush online as single command<commit_after><|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.

Send search terms as arguments to the command-line tool to search for
them. See the usage instructions. The persist value is set based on
an answer here: https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse
"""
import argparse
import os
import sys

# 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 tweets
from lib.twitter import auth, search

# 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):
    """
    Search Twitter API for tweets matching input search terms.

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

    @param searchQuery: search query as a string containing terms.
    @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.

    @return processedCount: count of tweets processed, as an integer.
        This will be a number up to the totalCount flag, but may less if
        fewer tweets are available in the 7-day window, or some tweets
        which were received were ignored because of language restriction
        we apply.
    """
    assert API_CONN, ("Authenticate with Twitter API before doing"
                      " a search for tweets.")

    searchRes = search.fetchTweetsPaging(API_CONN, searchQuery=searchQuery,
                                         itemLimit=totalCount)
    processedCount = 0
    for t in searchRes:
        if persist:
            # Add/update tweet author.
            profileRec = tweets.insertOrUpdateProfile(t.author)
            # Add/update the tweet.
            tweetData = tweets.insertOrUpdateTweet(t, profileRec.id)
        else:
            text = t.full_text.replace('\n', ' ').replace('\r', ' ')
            print u'@{0}: {1}'.format(t.author.screen_name, text)
        processedCount += 1

    return processedCount

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

    parser = argparse.ArgumentParser(
        description="Utility to search for tweets, then store tweet"
                    " and store profile data locally.",
        formatter_class=argparse.RawTextHelpFormatter
    )

    # Set value for terms list.
    searchQueryHelp = """\
Search tweet messages on Twitter API which match a query
string of one or more terms. See rules on the Twitter
API's search documentation.

Any double quotes or hashtags must be escaped on
the command-line, as covered below.

Examples:
 * message contains all terms, in any order.
    - wordA
    - wordA wordB wordC wordD
 * message contains at least one of the terms.
    - wordA OR wordB
    - to:handleA OR wordB
 * message contains hashtag terms.
    - \\#abc
    - '#def OR xyz OR #ghi'
 * message contains exact phrases. Note that
    quoted phrases must appear before other terms.
    - \\"My Quote\\"
    - '"My Quote" OR "Another quote" OR wordC'"""
    parser.add_argument('terms', metavar='TERM', nargs='+',
                        help=searchQueryHelp)

    # Set value count. Allows -c N or -c=N syntax.
    parser.add_argument('-c', '--count', metavar='N', type=int, default=200,
                        help="Default 200. Max count of tweets to get for the"
                             " search query.")

    # Set persist boolean value.
    parser.add_argument('--persist', dest='persist', action='store_true',
                        help="Store tweets and profiles in database."
                             " Defaults to on.")
    parser.add_argument('--no-persist', dest='persist', action='store_false',
                        help="Print tweet and profile data without storing.")
    parser.set_defaults(persist=True)

    args = parser.parse_args()

    # Combine trimmed list of strings into single string.
    searchQuery = u' '.join(args.terms)
    print u'Search query: {0}'.format(searchQuery)

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

    processedCount = searchAndStore(args.terms, totalCount=args.count,
                                    persist=args.persist)
    print 'Processed tweets: {0}'.format(processedCount)

if __name__ == '__main__':
    main()
<commit_msg>Update searchAndStoreTweets.py formatting<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.

Send search terms as arguments to the command-line tool to search for
them. See the usage instructions. The persist value is set based on
an answer here: https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse
"""
import argparse
import os
import sys

# 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 tweets
from lib.twitter import auth, search

# 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):
    """
    Search Twitter API for tweets matching input search terms.

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

    @param searchQuery: search query as a string containing terms.
    @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.

    @return processedCount: count of tweets processed, as an integer.
        This will be a number up to the totalCount flag, but may less if
        fewer tweets are available in the 7-day window, or some tweets
        which were received were ignored because of language restriction
        we apply.
    """
    assert API_CONN, ("Authenticate with Twitter API before doing"
                      " a search for tweets.")

    searchRes = search.fetchTweetsPaging(API_CONN, searchQuery=searchQuery,
                                         itemLimit=totalCount)
    processedCount = 0
    for t in searchRes:
        if persist:
            # Add/update tweet author.
            profileRec = tweets.insertOrUpdateProfile(t.author)
            # Add/update the tweet.
            tweetData = tweets.insertOrUpdateTweet(t, profileRec.id)
        else:
            text = t.full_text.replace('\n', ' ').replace('\r', ' ')
            print u'@{0}: {1}'.format(t.author.screen_name, text)
        processedCount += 1

    return processedCount

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

    parser = argparse.ArgumentParser(
        description="Utility to search for tweets then store tweet"
                    " and profile data locally.",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    # Set value for terms list.
    searchQueryHelp = """\
Search tweet messages on Twitter API which match a query
string of one or more terms. See rules on the Twitter
API's search documentation.

Any double quotes or hashtags must be escaped on
the command-line, as covered below.

Examples:
 * message contains all terms, in any order.
    - wordA
    - wordA wordB wordC wordD
 * message contains at least one of the terms.
    - wordA OR wordB
    - to:handleA OR wordB
 * message contains hashtag terms.
    - \\#abc
    - '#def OR xyz OR #ghi'
 * message contains exact phrases. Note that
    quoted phrases must appear before other terms.
    - \\"My Quote\\"
    - '"My Quote" OR "Another quote" OR wordC'"""
    parser.add_argument('terms', metavar='TERM', nargs='+',
                        help=searchQueryHelp)

    # Set value count. Allows -c N or -c=N syntax.
    parser.add_argument('-c', '--count', metavar='N', type=int, default=200,
                        help="Default 200. Max count of tweets to get for the"
                             " search query.")

    # Set persist boolean value.
    parser.add_argument('--persist', dest='persist', action='store_true',
                        help="Store tweets and profiles in database."
                             " Defaults to on.")
    parser.add_argument('--no-persist', dest='persist', action='store_false',
                        help="Print tweet and profile data without storing.")
    parser.set_defaults(persist=True)

    args = parser.parse_args()

    # Combine trimmed list of strings into single string.
    searchQuery = u" ".join(args.terms)
    print u'Search query: {0}'.format(searchQuery)

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

    processedCount = searchAndStore(args.terms, totalCount=args.count,
                                    persist=args.persist)
    print "Processed tweets: {0:,d}".format(processedCount)

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


<commit_before>bl_info = {
	"name": "Blender Queue",
	"author": "Patrick W. Crawford",
	"version": (1, 0),
	"blender": (2, 72, 1),
	"location": "Properties > Render",
	"description": "Render queues via command line",
	"warning": "",
	"wiki_url": "http://www.theduckcow.com",
	"category": "Render"}

import bpy, os, subprocess
from subprocess import Popen, PIPE

def renderQueue(self, context):

	print("hello render commandline")

	#get the blender path
	#subprocess.call(['open', '-W', '-a', 'Terminal.app', '&'])

	command = ['open', '-W', '-a', 'Terminal.app', '&']

	p = Popen(command, stdin=PIPE, stdout=PIPE, stderr=PIPE)
	stdout, err = p.communicate(b"")
	rc = p.returncode

class RENDER_queue(bpy.types.Operator):
	"""Command Line Render start"""
	bl_idname = "render.queue"
	bl_label = "Command Line"
	#bl_options = {'REGISTER', 'UNDO'}

	def execute(self, context):

		renderQueue(self, context)

		return {'FINISHED'}

# This allows you to right click on a button and link to the manual
"""
def add_object_manual_map():
	url_manual_prefix = "http://www.theduckcow.com"
	url_manual_mapping = (
		("bpy.ops.mesh.add_object", "Modeling/Objects"),
		)
	return url_manual_prefix, url_manual_mapping
"""

def render_panel(self, context):

	#layout = self.layout

	#add the two buttons here!
	self.layout.operator(
		RENDER_queue.bl_idname,
		text="Render Queue",
		icon='CONSOLE')

	#the list of queue'd objects, via text file
	"""
	ob = context.object
	group = ob.vertex_groups.active

	rows = 2
	if group:
		rows = 4
	row = layout.row()
	row.template_list("MESH_UL_vgroups", "", ob, "vertex_groups", ob.vertex_groups, "active_index", rows=rows)

	col = row.column(align=True)
	col.operator("object.vertex_group_add", icon='ZOOMIN', text="")
	col.operator("object.vertex_group_remove", icon='ZOOMOUT', text="").all = False
	col.menu("MESH_MT_vertex_group_specials", icon='DOWNARROW_HLT', text="")
	if group:
		col.separator()
		col.operator("object.vertex_group_move", icon='TRIA_UP', text="").direction = 'UP'
		col.operator("object.vertex_group_move", icon='TRIA_DOWN', text="").direction = 'DOWN'

	if ob.vertex_groups and (ob.mode == 'EDIT' or (ob.mode == 'WEIGHT_PAINT' and ob.type == 'MESH' and ob.data.use_paint_mask_vertex)):
		row = layout.row()

		sub = row.row(align=True)
		sub.operator("object.vertex_group_assign", text="Assign")
		sub.operator("object.vertex_group_remove_from", text="Remove")

		sub = row.row(align=True)
		sub.operator("object.vertex_group_select", text="Select")
		sub.operator("object.vertex_group_deselect", text="Deselect")

		layout.prop(context.tool_settings, "vertex_group_weight", text="Weight")
	"""

def register():
	bpy.utils.register_class(RENDER_queue)
	try:
		bpy.types.RENDER_PT_render.remove(render_panel)
	except:
		print("oh well")
	bpy.types.RENDER_PT_render.append(render_panel)

	#bpy.utils.register_manual_map(add_object_manual_map)
	#bpy.types.INFO_MT_mesh_add.append(add_object_button)

def unregister():
	bpy.utils.unregister_class(RENDER_queue)
	bpy.types.RENDER_PT_render.remove(render_panel)

	#bpy.utils.unregister_manual_map(add_object_manual_map)
	#bpy.types.INFO_MT_mesh_add.remove(add_object_button)

if __name__ == "__main__":
	register()
<commit_msg>just cleaning up and moving on<commit_after>bl_info = {
	"name": "Blender Queue",
	"author": "Patrick W. Crawford",
	"version": (1, 0),
	"blender": (2, 72, 1),
	"location": "Properties > Render",
	"description": "Render queues via command line",
	"warning": "",
	"wiki_url": "http://www.theduckcow.com",
	"category": "Render"}

import bpy, os, subprocess
from subprocess import Popen, PIPE

def renderQueue(self, context):

	print("hello render commandline")

	blendPath = bpy.path.abspath("//") # could also do bpy.data.filepath
	print("-b {x} -x 1 -a".format(x=blendPath))

	""" I have no idea what I'm doing (below) """

	#get the blender path

	#below shows currently open command line OR opens new if not already open
	command = ['open', '-W', '-a', 'Terminal.app']
	# and blender would return and be active once you closed the terminal (..twice. first time closed reappears)

	#command = ['open', '-W', '-a', 'Terminal.app', '~']
	#makes it do nothing. Through just terminal, will launch new window regardless at that directory
	# even adding "&" made it work as I wanted!!!!!! ie in terminal: "open -W -a Terminal.app ~/Desktop &"
	# opened a new terminal window at said directory, but original process would keep going

	#subprocess.call(command)
	#below won't freeze blender, but then can't see shell
	#subprocess.Popen(command, creationflags=0x00000008)
	# yay! windows only -_-

	# will NOT freeze blender! popen is the proper way to do background...
	# but same deal, if terminal already open just brings that forward..

	#os.system("start cmd /K dir")

	# below works on command line (mac at least), but not through python so far
	#command = ['open', '-a', 'Terminal', "`pwd`"]
	#command = []
	"""
	p = Popen(command, stdin=PIPE, stdout=PIPE, stderr=PIPE)
	stdout, err = p.communicate(b"")
	rc = p.returncode
	"""

class RENDER_queue(bpy.types.Operator):
	"""Command Line Render start"""
	bl_idname = "render.queue"
	bl_label = "Command Line"
	#bl_options = {'REGISTER', 'UNDO'}

	def execute(self, context):

		renderQueue(self, context)

		return {'FINISHED'}

# This allows you to right click on a button and link to the manual
"""
def add_object_manual_map():
	url_manual_prefix = "http://www.theduckcow.com"
	url_manual_mapping = (
		("bpy.ops.mesh.add_object", "Modeling/Objects"),
		)
	return url_manual_prefix, url_manual_mapping
"""

def render_panel(self, context):

	#layout = self.layout

	#add the two buttons here!
	self.layout.operator(
		RENDER_queue.bl_idname,
		text="Render Queue",
		icon='CONSOLE')

	#the list of queue'd objects, via text file
	"""
	ob = context.object
	group = ob.vertex_groups.active

	rows = 2
	if group:
		rows = 4
	row = layout.row()
	row.template_list("MESH_UL_vgroups", "", ob, "vertex_groups", ob.vertex_groups, "active_index", rows=rows)

	col = row.column(align=True)
	col.operator("object.vertex_group_add", icon='ZOOMIN', text="")
	col.operator("object.vertex_group_remove", icon='ZOOMOUT', text="").all = False
	col.menu("MESH_MT_vertex_group_specials", icon='DOWNARROW_HLT', text="")
	if group:
		col.separator()
		col.operator("object.vertex_group_move", icon='TRIA_UP', text="").direction = 'UP'
		col.operator("object.vertex_group_move", icon='TRIA_DOWN', text="").direction = 'DOWN'

	if ob.vertex_groups and (ob.mode == 'EDIT' or (ob.mode == 'WEIGHT_PAINT' and ob.type == 'MESH' and ob.data.use_paint_mask_vertex)):
		row = layout.row()

		sub = row.row(align=True)
		sub.operator("object.vertex_group_assign", text="Assign")
		sub.operator("object.vertex_group_remove_from", text="Remove")

		sub = row.row(align=True)
		sub.operator("object.vertex_group_select", text="Select")
		sub.operator("object.vertex_group_deselect", text="Deselect")

		layout.prop(context.tool_settings, "vertex_group_weight", text="Weight")
	"""

def register():
	bpy.utils.register_class(RENDER_queue)
	try:
		bpy.types.RENDER_PT_render.remove(render_panel)
	except:
		print("oh well")
	bpy.types.RENDER_PT_render.append(render_panel)

	#bpy.utils.register_manual_map(add_object_manual_map)
	#bpy.types.INFO_MT_mesh_add.append(add_object_button)

def unregister():
	bpy.utils.unregister_class(RENDER_queue)
	bpy.types.RENDER_PT_render.remove(render_panel)

	#bpy.utils.unregister_manual_map(add_object_manual_map)
	#bpy.types.INFO_MT_mesh_add.remove(add_object_button)

if __name__ == "__main__":
	register()
<|endoftext|>


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

import unittest
from har2tree import CrawledTree
from pathlib import Path
import os
import uuid

class SimpleTest(unittest.TestCase):

    index_tree: CrawledTree

    @classmethod
    def setUpClass(cls) -> None:
        test_dir = Path(os.path.abspath(os.path.dirname(__file__))) / 'data' / 'simple'
        har_to_process = [test_dir / 'heroku_index.har']
        cls.index_tree = CrawledTree(har_to_process, str(uuid.uuid4()))

    def test_root_url(self) -> None:
        self.assertEqual(self.index_tree.root_url, 'https://lookyloo-testing.herokuapp.com/')

    def test_user_agent(self) -> None:
        self.assertEqual(self.index_tree.user_agent, "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/89.0.4389.128 Safari/537.36")

if __name__ == '__main__':
    unittest.main()
<commit_msg>chg: tests now work with submodule capture & added redirect test<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import unittest
from har2tree import CrawledTree
from pathlib import Path
import os
import uuid

class SimpleTest(unittest.TestCase):

    http_redirect_tree: CrawledTree

    @classmethod
    def setUpClass(cls) -> None:
        test_dir = Path(os.path.abspath(os.path.dirname(__file__))) / 'capture_samples' / 'http_redirect'
        har_to_process = [test_dir / '0.har']
        cls.http_redirect_tree = CrawledTree(har_to_process, str(uuid.uuid4()))

    def test_root_url(self) -> None:
        self.assertEqual(self.http_redirect_tree.root_url, 'https://lookyloo-testing.herokuapp.com/redirect_http')

    def test_user_agent(self) -> None:
        self.assertEqual(self.http_redirect_tree.user_agent, "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/534+ (KHTML, like Gecko) BingPreview/1.0b")

    def test_redirects(self) -> None:
        self.assertEqual(self.http_redirect_tree.redirects[1], "https://www.youtube.com/watch?v=iwGFalTRHDA")

if __name__ == '__main__':
    unittest.main()

<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
# guppi.py

A python file handler for guppi RAW files from the GBT.

The guppi raw format consists of a FITS-like header, followed by a block of data,
and repeated over and over until the end of the file.
"""

import numpy as np
import os
import time
from pprint import pprint
from astropy.coordinates import Angle

from .utils import unpack, rebin
import sys

PYTHON3 = sys.version_info >= (3, 0)

# Check if $DISPLAY is set (for handling plotting on remote machines with no X-forwarding)
if 'DISPLAY' in os.environ.keys():

    try:
        import matplotlib
        #matplotlib.use('Qt5Agg')
    except ImportError:
        pass
    import pylab as plt
else:
    import matplotlib
    matplotlib.use('Agg')
    import pylab as plt

###
# Config values
###

MAX_PLT_POINTS      = 65536 * 4              # Max number of points in matplotlib plot
MAX_IMSHOW_POINTS   = (8192, 4096)           # Max number of points in imshow plot
MAX_DATA_ARRAY_SIZE = 1024 * 1024 * 1024     # Max size of data array to load into memory

class EndOfFileError(Exception):
    pass

class GuppiRaw(object):
    """ Python class for reading Guppi raw files

    Args:
        filename (str): name of the .raw file to open

    Optional args:
        n_blocks (int): if number of blocks to read is known, set it here.
                        This saves seeking through the file to check how many
                        integrations there are in the file.
    """
    def __init__(self, filename, n_blocks=None):
        self.filename = filename
        if PYTHON3:
            self.file_obj = open(filename, 'rb')
        else:
            self.file_obj = open(filename, 'r')
        self.filesize = os.path.getsize(filename)

        if not n_blocks:
            self.n_blocks = self.find_n_data_blocks()

        else:
            self.n_blocks = n_blocks

        self._d = np.zeros(1, dtype='complex64')
        self._d_x = np.zeros(1, dtype='int8')
        self._d_y = np.zeros(1, dtype='int8')

    def __enter__(self):
         return self

    def __exit__(self, exception_type, exception_value, traceback):
        self.file_obj.close()

    def __repr__(self):
        return "<GuppiRaw file handler for %s>" % self.filename

    def read_header(self):
        """ Read next header (multiple headers in file)

        Returns:
            (header, data_idx) - a dictionary of keyword:value header data and
            also the byte index of where the corresponding data block resides.
        """
        start_idx = self.file_obj.tell()
        key, val = '', ''

        header_dict = {}
        keep_reading = True

        first_line = self.file_obj

        try:
            while keep_reading:
                if start_idx + 80 > self.filesize:
                    keep_reading = False
                    raise EndOfFileError
                line = self.file_obj.read(80)
                if PYTHON3:
                    line = line.decode("utf-8")
                #print line
                if line.startswith('END'):
                    keep_reading = False
                    break
                else:
                    key, val = line.split('=')
                    key, val = key.strip(), val.strip()

                    if "'" in val:
                        # Items in quotes are strings
                        val = str(val.strip("'").strip())
                    elif "." in val:
                        # Items with periods are floats (if not a string)
                        val = float(val)
                    else:
                        # Otherwise it's an integer
                        val = int(val)

                header_dict[key] = val
        except ValueError:
            print("CURRENT LINE: ", line)
            print("BLOCK START IDX: ",  start_idx)
            print("FILE SIZE: ",  self.filesize)
            print("NEXT 512 BYTES: \n")
            print(self.file_obj.read(512))
            raise

        data_idx = self.file_obj.tell()

        # Seek past padding if DIRECTIO is being used
        if "DIRECTIO" in header_dict.keys():
            if int(header_dict["DIRECTIO"]) == 1:
                if data_idx % 512:
                    data_idx += (512 - data_idx % 512)

        self.file_obj.seek(start_idx)
        return header_dict, data_idx

    def read_first_header(self):
        """ Read first header in file

        Returns:
            header (dict): keyword:value pairs of header metadata
        """
        self.file_obj.seek(0)
        header_dict, pos = self.read_header()
        self.file_obj.seek(0)
        return header_dict

    def read_next_data_block_shape(self):
        header, data_idx = self.read_header()
        n_chan = int(header['OBSNCHAN'])
        n_pol  = min(2, int(header['NPOL']))  #NPOL 4 actually means 2
        n_samples = int(header['BLOCSIZE']) / n_chan / n_pol

        is_chanmaj = False
        if 'CHANMAJ' in header.keys():
            if int(header['CHANMAJ']) == 1:
                is_chanmaj = True
        if is_chanmaj:
            dshape = (int(n_samples), n_chan, n_pol)
        else:
            dshape = (n_chan, int(n_samples), n_pol)
        return dshape

    def read_next_data_block_int8(self):
        """ Read the next block of data and its header

        Returns: (header, data)
            header (dict): dictionary of header metadata
            data (np.array): Numpy array of data, converted into to complex64.
        """
        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)

        # Read data and reshape

        n_chan = int(header['OBSNCHAN'])
        n_pol  = int(header['NPOL'])
        n_samples = int(header['BLOCSIZE']) / n_chan / n_pol
        n_bit = int(header['NBITS'])

        d = np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8')

        # Handle 2-bit and 4-bit data
        if n_bit != 8:
            d = unpack(d, n_bit)

        d = d.reshape((n_chan, n_samples, n_pol))    # Real, imag

        if self._d_x.shape != d[..., 0:2].shape:
                self._d_x = np.ascontiguousarray(np.zeros(d[..., 0:2].shape, dtype='int8'))
                self._d_y = np.ascontiguousarray(np.zeros(d[..., 2:4].shape, dtype='int8'))

        self._d_x[:] = d[..., 0:2]
        self._d_y[:] = d[..., 2:4]
        return header, self._d_x, self._d_y

    def read_next_data_block_int8_2x(self):
        """ Read the next block of data and its header

        Returns: (header, data)
            header (dict): dictionary of header metadata
            data (np.array): Numpy array of data, converted into to complex64.
        """
        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)

        # Read data and reshape

        n_chan = int(header['OBSNCHAN'])
        n_pol  = int(header['NPOL'])
        n_samples = int(header['BLOCSIZE']) / n_chan / n_pol
        n_bit = int(header['NBITS'])

        d = np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8')

        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)
        d2 = np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8')

        # Handle 2-bit and 4-bit data
        if n_bit != 8:
            d = unpack(d, n_bit)

        d = d.reshape((n_chan, n_samples, n_pol))    # Real, imag
        d2 = d2.reshape((n_chan, n_samples, n_pol))
        d = np.concatenate((d, d2), axis=1)
        print(d.shape)

        if self._d_x.shape != (n_chan, n_samples*2, n_pol):
                self._d_x = np.ascontiguousarray(np.zeros(d[..., 0:2].shape, dtype='int8'))
                self._d_y = np.ascontiguousarray(np.zeros(d[..., 2:4].shape, dtype='int8'))

        self._d_x[:] = d[..., 0:2]
        self._d_y[:] = d[..., 2:4]
        return header, self._d_x, self._d_y

    def read_next_data_block(self):
        """ Read the next block of data and its header

        Returns: (header, data)
            header (dict): dictionary of header metadata
            data (np.array): Numpy array of data, converted into to complex64.
        """
        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)

        # Read data and reshape

        n_chan = int(header['OBSNCHAN'])
        n_pol  = int(header['NPOL'])
        n_samples = int(header['BLOCSIZE']) / n_chan / n_pol
        n_bit = int(header['NBITS'])

        d = np.ascontiguousarray(np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8'))

        # Handle 2-bit and 4-bit data
        if n_bit != 8:
            d = unpack(d, n_bit)

        dshape = self.read_next_data_block_shape()

        d = d.reshape(dshape)    # Real, imag

        if self._d.shape != d.shape:
            self._d = np.zeros(d.shape, dtype='float32')

        self._d[:] = d

        return header, self._d[:].view('complex64')

    def find_n_data_blocks(self):
        """ Seek through the file to find how many data blocks there are in the file

        Returns:
            n_blocks (int): number of data blocks in the file
        """
        self.file_obj.seek(0)
        header0, data_idx0 = self.read_header()

        self.file_obj.seek(data_idx0)
        block_size = int(header0['BLOCSIZE'])
        n_bits     = int(header0['NBITS'])
        self.file_obj.seek(int(header0['BLOCSIZE']), 1)
        n_blocks = 1
        end_found = False
        while not end_found:
            try:
                header, data_idx = self.read_header()
                self.file_obj.seek(data_idx)
                self.file_obj.seek(header['BLOCSIZE'], 1)
                n_blocks += 1
            except EndOfFileError:
                end_found = True
                break

        self.file_obj.seek(0)
        return n_blocks

    def reset_index(self):
        """ Return file_obj seek to start of file """
        self.file_obj.seek(0)

    def print_stats(self):
        """ Compute some basic stats on the next block of data """

        header, data = self.read_next_data_block()
        data = data.view('float32')

        print("AVG: %2.3f" % data.mean())
        print("STD: %2.3f" % data.std())
        print("MAX: %2.3f" % data.max())
        print("MIN: %2.3f" % data.min())

        import pylab as plt

    def plot_histogram(self, filename=None):
        """ Plot a histogram of data values """
        header, data = self.read_next_data_block()
        data = data.view('float32')

        plt.figure("Histogram")
        plt.hist(data.flatten(), 65, facecolor='#cc0000')
        if filename:
            plt.savefig(filename)
        plt.show()

    def plot_spectrum(self, filename=None, plot_db=True):
        """ Do a (slow) numpy FFT and take power of data """
        header, data = self.read_next_data_block()

        print("Computing FFT...")
        d_xx_fft = np.abs(np.fft.fft(data[..., 0]))
        d_xx_fft = d_xx_fft.flatten()

        # Rebin to max number of points
        dec_fac_x = 1
        if d_xx_fft.shape[0] > MAX_PLT_POINTS:
            dec_fac_x = d_xx_fft.shape[0] / MAX_PLT_POINTS

        d_xx_fft  = rebin(d_xx_fft, dec_fac_x)

        print("Plotting...")
        if plot_db:
            plt.plot(10 * np.log10(d_xx_fft))
            plt.ylabel("Power [dB]")
        else:
            plt.plot(d_xx_fft)
            plt.ylabel("Power")
            plt.xlabel("Channel")
        plt.title(self.filename)
        if filename:
            plt.savefig(filename)
        plt.show()

    def generate_filterbank_header(self, nchans=1, ):
        """ Generate a blimpy header dictionary """
        gp_head = self.read_first_header()
        fb_head = {}

        telescope_str = gp_head.get("TELESCOP", "unknown")
        if telescope_str in ('GBT', 'GREENBANK'):
            fb_head["telescope_id"] = 6
        elif telescope_str in ('PKS', 'PARKES'):
            fb_head["telescop_id"] = 7
        else:
            fb_head["telescop_id"] = 0

        # Using .get() method allows us to fill in default values if not present
        fb_head["source_name"] = gp_head.get("SRC_NAME", "unknown")
        fb_head["az_start"]    = gp_head.get("AZ", 0)
        fb_head["za_start"]    = gp_head.get("ZA", 0)

        fb_head["src_raj"]     = Angle(str(gp_head.get("RA", 0.0)) + "hr")
        fb_head["src_dej"]     = Angle(str(gp_head.get("DEC", 0.0)) + "deg")
        fb_head["rawdatafile"] = self.filename

        # hardcoded
        fb_head["machine_id"]    = 20
        fb_head["data_type"]     = 1      # blio datatype
        fb_head["barycentric"]      = 0
        fb_head["pulsarcentric"] = 0
        fb_head["nbits"]         = 32

        # TODO - compute these values. Need to figure out the correct calcs
        fb_head["tstart"]        = 0.0
        fb_head["tsamp"]         = 1.0
        fb_head["fch1"]          = 0.0
        fb_head["foff"]          = 187.5 / nchans

        # Need to be updated based on output specs
        fb_head["nchans"]        = nchans
        fb_head["nifs"]          = 1
        fb_head["nbeams"]        = 1

        return fb_head

def cmd_tool(args=None):
    """ Command line tool for plotting and viewing info on guppi raw files """

    from argparse import ArgumentParser

    parser = ArgumentParser(description="Command line utility for creating spectra from GuppiRaw files.")

    parser.add_argument('filename', type=str, help='Name of file to read')
    parser.add_argument('-o', dest='outdir', type=str, default='./', help='output directory for PNG files')
    args = parser.parse_args()

    r = GuppiRaw(args.filename)

    r.print_stats()
    bname = os.path.splitext(os.path.basename(args.filename))[0]
    bname = os.path.join(args.outdir, bname)
    r.plot_histogram(filename="%s_hist.png" % bname)
    r.plot_spectrum(filename="%s_spec.png" % bname)

if __name__ == "__main__":
    cmd_tool()
<commit_msg>revert last commit<commit_after>#!/usr/bin/env python
"""
# guppi.py

A python file handler for guppi RAW files from the GBT.

The guppi raw format consists of a FITS-like header, followed by a block of data,
and repeated over and over until the end of the file.
"""

import numpy as np
import os
import time
from pprint import pprint
from astropy.coordinates import Angle

from .utils import unpack, rebin
import sys

PYTHON3 = sys.version_info >= (3, 0)

# Check if $DISPLAY is set (for handling plotting on remote machines with no X-forwarding)
if 'DISPLAY' in os.environ.keys():

    try:
        import matplotlib
        #matplotlib.use('Qt5Agg')
    except ImportError:
        pass
    import pylab as plt
else:
    import matplotlib
    matplotlib.use('Agg')
    import pylab as plt

###
# Config values
###

MAX_PLT_POINTS      = 65536 * 4              # Max number of points in matplotlib plot
MAX_IMSHOW_POINTS   = (8192, 4096)           # Max number of points in imshow plot
MAX_DATA_ARRAY_SIZE = 1024 * 1024 * 1024     # Max size of data array to load into memory

class EndOfFileError(Exception):
    pass

class GuppiRaw(object):
    """ Python class for reading Guppi raw files

    Args:
        filename (str): name of the .raw file to open

    Optional args:
        n_blocks (int): if number of blocks to read is known, set it here.
                        This saves seeking through the file to check how many
                        integrations there are in the file.
    """
    def __init__(self, filename, n_blocks=None):
        self.filename = filename
        if PYTHON3:
            self.file_obj = open(filename, 'rb')
        else:
            self.file_obj = open(filename, 'r')
        self.filesize = os.path.getsize(filename)

        if not n_blocks:
            self.n_blocks = self.find_n_data_blocks()

        else:
            self.n_blocks = n_blocks

        self._d = np.zeros(1, dtype='complex64')
        self._d_x = np.zeros(1, dtype='int8')
        self._d_y = np.zeros(1, dtype='int8')

    def __enter__(self):
         return self

    def __exit__(self, exception_type, exception_value, traceback):
        self.file_obj.close()

    def __repr__(self):
        return "<GuppiRaw file handler for %s>" % self.filename

    def read_header(self):
        """ Read next header (multiple headers in file)

        Returns:
            (header, data_idx) - a dictionary of keyword:value header data and
            also the byte index of where the corresponding data block resides.
        """
        start_idx = self.file_obj.tell()
        key, val = '', ''

        header_dict = {}
        keep_reading = True

        first_line = self.file_obj

        try:
            while keep_reading:
                if start_idx + 80 > self.filesize:
                    keep_reading = False
                    raise EndOfFileError
                line = self.file_obj.read(80)
                if PYTHON3:
                    line = line.decode("utf-8")
                #print line
                if line.startswith('END'):
                    keep_reading = False
                    break
                else:
                    key, val = line.split('=')
                    key, val = key.strip(), val.strip()

                    if "'" in val:
                        # Items in quotes are strings
                        val = str(val.strip("'").strip())
                    elif "." in val:
                        # Items with periods are floats (if not a string)
                        val = float(val)
                    else:
                        # Otherwise it's an integer
                        val = int(val)

                header_dict[key] = val
        except ValueError:
            print("CURRENT LINE: ", line)
            print("BLOCK START IDX: ",  start_idx)
            print("FILE SIZE: ",  self.filesize)
            print("NEXT 512 BYTES: \n")
            print(self.file_obj.read(512))
            raise

        data_idx = self.file_obj.tell()

        # Seek past padding if DIRECTIO is being used
        if "DIRECTIO" in header_dict.keys():
            if int(header_dict["DIRECTIO"]) == 1:
                if data_idx % 512:
                    data_idx += (512 - data_idx % 512)

        self.file_obj.seek(start_idx)
        return header_dict, data_idx

    def read_first_header(self):
        """ Read first header in file

        Returns:
            header (dict): keyword:value pairs of header metadata
        """
        self.file_obj.seek(0)
        header_dict, pos = self.read_header()
        self.file_obj.seek(0)
        return header_dict

    def read_next_data_block_shape(self):
        header, data_idx = self.read_header()
        n_chan = int(header['OBSNCHAN'])
        n_pol  = int(header['NPOL'])
        n_samples = int(header['BLOCSIZE']) / n_chan / n_pol

        is_chanmaj = False
        if 'CHANMAJ' in header.keys():
            if int(header['CHANMAJ']) == 1:
                is_chanmaj = True
        if is_chanmaj:
            dshape = (int(n_samples), n_chan, n_pol)
        else:
            dshape = (n_chan, int(n_samples), n_pol)
        return dshape

    def read_next_data_block_int8(self):
        """ Read the next block of data and its header

        Returns: (header, data)
            header (dict): dictionary of header metadata
            data (np.array): Numpy array of data, converted into to complex64.
        """
        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)

        # Read data and reshape

        n_chan = int(header['OBSNCHAN'])
        n_pol  = int(header['NPOL'])
        n_samples = int(header['BLOCSIZE']) / n_chan / n_pol
        n_bit = int(header['NBITS'])

        d = np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8')

        # Handle 2-bit and 4-bit data
        if n_bit != 8:
            d = unpack(d, n_bit)

        d = d.reshape((n_chan, n_samples, n_pol))    # Real, imag

        if self._d_x.shape != d[..., 0:2].shape:
                self._d_x = np.ascontiguousarray(np.zeros(d[..., 0:2].shape, dtype='int8'))
                self._d_y = np.ascontiguousarray(np.zeros(d[..., 2:4].shape, dtype='int8'))

        self._d_x[:] = d[..., 0:2]
        self._d_y[:] = d[..., 2:4]
        return header, self._d_x, self._d_y

    def read_next_data_block_int8_2x(self):
        """ Read the next block of data and its header

        Returns: (header, data)
            header (dict): dictionary of header metadata
            data (np.array): Numpy array of data, converted into to complex64.
        """
        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)

        # Read data and reshape

        n_chan = int(header['OBSNCHAN'])
        n_pol  = int(header['NPOL'])
        n_samples = int(header['BLOCSIZE']) / n_chan / n_pol
        n_bit = int(header['NBITS'])

        d = np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8')

        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)
        d2 = np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8')

        # Handle 2-bit and 4-bit data
        if n_bit != 8:
            d = unpack(d, n_bit)

        d = d.reshape((n_chan, n_samples, n_pol))    # Real, imag
        d2 = d2.reshape((n_chan, n_samples, n_pol))
        d = np.concatenate((d, d2), axis=1)
        print(d.shape)

        if self._d_x.shape != (n_chan, n_samples*2, n_pol):
                self._d_x = np.ascontiguousarray(np.zeros(d[..., 0:2].shape, dtype='int8'))
                self._d_y = np.ascontiguousarray(np.zeros(d[..., 2:4].shape, dtype='int8'))

        self._d_x[:] = d[..., 0:2]
        self._d_y[:] = d[..., 2:4]
        return header, self._d_x, self._d_y

    def read_next_data_block(self):
        """ Read the next block of data and its header

        Returns: (header, data)
            header (dict): dictionary of header metadata
            data (np.array): Numpy array of data, converted into to complex64.
        """
        header, data_idx = self.read_header()
        self.file_obj.seek(data_idx)

        # Read data and reshape

        n_chan = int(header['OBSNCHAN'])
        n_pol  = int(header['NPOL'])
        n_samples = int(header['BLOCSIZE']) / n_chan / n_pol
        n_bit = int(header['NBITS'])

        d = np.ascontiguousarray(np.fromfile(self.file_obj, count=header['BLOCSIZE'], dtype='int8'))

        # Handle 2-bit and 4-bit data
        if n_bit != 8:
            d = unpack(d, n_bit)

        dshape = self.read_next_data_block_shape()

        d = d.reshape(dshape)    # Real, imag

        if self._d.shape != d.shape:
            self._d = np.zeros(d.shape, dtype='float32')

        self._d[:] = d

        return header, self._d[:].view('complex64')

    def find_n_data_blocks(self):
        """ Seek through the file to find how many data blocks there are in the file

        Returns:
            n_blocks (int): number of data blocks in the file
        """
        self.file_obj.seek(0)
        header0, data_idx0 = self.read_header()

        self.file_obj.seek(data_idx0)
        block_size = int(header0['BLOCSIZE'])
        n_bits     = int(header0['NBITS'])
        self.file_obj.seek(int(header0['BLOCSIZE']), 1)
        n_blocks = 1
        end_found = False
        while not end_found:
            try:
                header, data_idx = self.read_header()
                self.file_obj.seek(data_idx)
                self.file_obj.seek(header['BLOCSIZE'], 1)
                n_blocks += 1
            except EndOfFileError:
                end_found = True
                break

        self.file_obj.seek(0)
        return n_blocks

    def reset_index(self):
        """ Return file_obj seek to start of file """
        self.file_obj.seek(0)

    def print_stats(self):
        """ Compute some basic stats on the next block of data """

        header, data = self.read_next_data_block()
        data = data.view('float32')

        print("AVG: %2.3f" % data.mean())
        print("STD: %2.3f" % data.std())
        print("MAX: %2.3f" % data.max())
        print("MIN: %2.3f" % data.min())

        import pylab as plt

    def plot_histogram(self, filename=None):
        """ Plot a histogram of data values """
        header, data = self.read_next_data_block()
        data = data.view('float32')

        plt.figure("Histogram")
        plt.hist(data.flatten(), 65, facecolor='#cc0000')
        if filename:
            plt.savefig(filename)
        plt.show()

    def plot_spectrum(self, filename=None, plot_db=True):
        """ Do a (slow) numpy FFT and take power of data """
        header, data = self.read_next_data_block()

        print("Computing FFT...")
        d_xx_fft = np.abs(np.fft.fft(data[..., 0]))
        d_xx_fft = d_xx_fft.flatten()

        # Rebin to max number of points
        dec_fac_x = 1
        if d_xx_fft.shape[0] > MAX_PLT_POINTS:
            dec_fac_x = d_xx_fft.shape[0] / MAX_PLT_POINTS

        d_xx_fft  = rebin(d_xx_fft, dec_fac_x)

        print("Plotting...")
        if plot_db:
            plt.plot(10 * np.log10(d_xx_fft))
            plt.ylabel("Power [dB]")
        else:
            plt.plot(d_xx_fft)
            plt.ylabel("Power")
            plt.xlabel("Channel")
        plt.title(self.filename)
        if filename:
            plt.savefig(filename)
        plt.show()

    def generate_filterbank_header(self, nchans=1, ):
        """ Generate a blimpy header dictionary """
        gp_head = self.read_first_header()
        fb_head = {}

        telescope_str = gp_head.get("TELESCOP", "unknown")
        if telescope_str in ('GBT', 'GREENBANK'):
            fb_head["telescope_id"] = 6
        elif telescope_str in ('PKS', 'PARKES'):
            fb_head["telescop_id"] = 7
        else:
            fb_head["telescop_id"] = 0

        # Using .get() method allows us to fill in default values if not present
        fb_head["source_name"] = gp_head.get("SRC_NAME", "unknown")
        fb_head["az_start"]    = gp_head.get("AZ", 0)
        fb_head["za_start"]    = gp_head.get("ZA", 0)

        fb_head["src_raj"]     = Angle(str(gp_head.get("RA", 0.0)) + "hr")
        fb_head["src_dej"]     = Angle(str(gp_head.get("DEC", 0.0)) + "deg")
        fb_head["rawdatafile"] = self.filename

        # hardcoded
        fb_head["machine_id"]    = 20
        fb_head["data_type"]     = 1      # blio datatype
        fb_head["barycentric"]      = 0
        fb_head["pulsarcentric"] = 0
        fb_head["nbits"]         = 32

        # TODO - compute these values. Need to figure out the correct calcs
        fb_head["tstart"]        = 0.0
        fb_head["tsamp"]         = 1.0
        fb_head["fch1"]          = 0.0
        fb_head["foff"]          = 187.5 / nchans

        # Need to be updated based on output specs
        fb_head["nchans"]        = nchans
        fb_head["nifs"]          = 1
        fb_head["nbeams"]        = 1

        return fb_head

def cmd_tool(args=None):
    """ Command line tool for plotting and viewing info on guppi raw files """

    from argparse import ArgumentParser

    parser = ArgumentParser(description="Command line utility for creating spectra from GuppiRaw files.")

    parser.add_argument('filename', type=str, help='Name of file to read')
    parser.add_argument('-o', dest='outdir', type=str, default='./', help='output directory for PNG files')
    args = parser.parse_args()

    r = GuppiRaw(args.filename)

    r.print_stats()
    bname = os.path.splitext(os.path.basename(args.filename))[0]
    bname = os.path.join(args.outdir, bname)
    r.plot_histogram(filename="%s_hist.png" % bname)
    r.plot_spectrum(filename="%s_spec.png" % bname)

if __name__ == "__main__":
    cmd_tool()
<|endoftext|>