<commit_before>from flask import current_app, jsonify, render_template
import json
import os
import re
import requests
from html.parser import HTMLParser

from na_common.dates import get_nice_event_dates

from app.comms.encryption import encrypt
from app.errors import InvalidRequest
from app.models import BASIC, EVENT, MAGAZINE
from app.dao.events_dao import dao_get_event_by_id
from app.dao.emails_dao import dao_get_past_hour_email_count_for_provider, dao_get_todays_email_count_for_provider
from app.dao.email_providers_dao import dao_get_first_email_provider, dao_get_next_email_provider
from app.dao.magazines_dao import dao_get_magazine_by_id

h = HTMLParser()

def get_email_provider(override=False, email_provider=None):
    if not email_provider:
        email_provider = dao_get_first_email_provider()
        if not email_provider:
            return None

    daily_email_count = dao_get_todays_email_count_for_provider(email_provider.id)
    hourly_email_count = 0

    if email_provider.hourly_limit > 0:
        hourly_email_count = dao_get_past_hour_email_count_for_provider(email_provider.id)

    if daily_email_count > email_provider.daily_limit:
        if override:
            next_email_provider = dao_get_next_email_provider(email_provider.pos)

            if next_email_provider:
                return get_email_provider(override, next_email_provider)
        else:
            email_provider.limit = 0
            email_provider.daily_limit_reached = True
        return email_provider

    if hourly_email_count > email_provider.hourly_limit:
        if override:
            next_email_provider = dao_get_next_email_provider(email_provider.pos)
            if next_email_provider:
                return get_email_provider(override, next_email_provider)
        email_provider.limit = 0
        email_provider.hourly_limit_reached = True

    if email_provider.hourly_limit > 0:
        email_provider.limit = email_provider.hourly_limit - hourly_email_count

    email_provider.limit = email_provider.daily_limit - daily_email_count
    return email_provider

def get_email_html(email_type, **kwargs):
    if email_type == EVENT:
        event = dao_get_event_by_id(kwargs.get('event_id'))
        member_id = kwargs.get('member_id')
        if not member_id:
            member_id = '0'  # for preview of emails
        current_app.logger.debug('Email Tokens %s', current_app.config['EMAIL_TOKENS'])
        unsubcode = encrypt(
            "{}={}".format(current_app.config['EMAIL_TOKENS']['member_id'], member_id),
            current_app.config['EMAIL_UNSUB_SALT']
        )
        return render_template(
            'emails/events.html',
            event=event,
            event_dates=get_nice_event_dates(event.event_dates),
            description=h.unescape(event.description),
            details=kwargs.get('details'),
            extra_txt=kwargs.get('extra_txt'),
            unsubcode=unsubcode
        )
    elif email_type == MAGAZINE:
        magazine = dao_get_magazine_by_id(kwargs.get('magazine_id'))
        topics = []
        if magazine.topics:
            _topics = [(t.split(':')[0], t.split(':')[1]) for t in magazine.topics.split('\n')]

            for title, description in _topics:
                topics.append({'title': title, 'description': description})
        return render_template(
            'emails/magazine.html',
            magazine=magazine,
            topics=topics
        )
    elif email_type == BASIC:
        return render_template(
            'emails/basic.html',
            title=kwargs.get('title', ''),
            message=kwargs.get('message')
        )

def get_email_data(data_map, to, subject, message, from_email, from_name):
    data_struct = {}
    if isinstance(to, list):
        to = ','.join(to)

    def set_data(_map, val):
        elements = _map.split(",")

        if len(elements) == 2:
            match = re.match(r"\[(?P<list_element>.+)\]", elements[1])
            if match:
                data_struct[elements[0]] = []

                for item in val.split(','):
                    data_struct[elements[0]].append({
                        match.group('list_element'): item
                    })
            else:
                if data_struct.get(elements[0]):
                    data_struct[elements[0]][elements[1]] = val
                else:
                    data_struct[elements[0]] = {elements[1]: val}
        else:
            data_struct[elements[0]] = val

    set_data(data_map['to'], to)
    set_data(data_map['subject'], subject)
    set_data(data_map['from'], from_email)
    if 'from_name' in data_map.keys():
        set_data(data_map['from_name'], from_name)
    set_data(data_map['message'], message)

    return data_struct

def send_email(to, subject, message, from_email=None, from_name=None, override=False):
    if current_app.config['EMAIL_DISABLED']:
        current_app.logger.info("Emails disabled, unset EMAIL_DISABLED env var to re-enable")
        return 200
    if current_app.config['ENVIRONMENT'] != 'live' or current_app.config.get('EMAIL_RESTRICT'):
        message = message.replace('', '
Test email, intended for {}
'.format(to))
        to = current_app.config['TEST_EMAIL']

    if not from_email:
        from_email = 'noreply@{}'.format(current_app.config['EMAIL_DOMAIN'])
    if not from_name:
        from_name = 'New Acropolis'

    email_provider = get_email_provider(override)

    if email_provider:
        if hasattr(email_provider, "hourly_limit_reached"):
            raise InvalidRequest('Hourly limit reached', 429)
        elif hasattr(email_provider, "daily_limit_reached"):
            raise InvalidRequest('Daily limit reached', 429)

        data = get_email_data(email_provider.data_map, to, subject, message, from_email, from_name)
        data = data if email_provider.as_json else json.dumps(data)

        headers = {
            'api-key': email_provider.api_key,
            'accept': 'application/json',
            'content-type': 'application/json'
        } if email_provider.headers else None

        response = requests.post(
            email_provider.api_url,
            auth=('api', email_provider.api_key),
            headers=headers,
            data=data,
        )

        response.raise_for_status()
        current_app.logger.info('Sent email: {}, response: {}'.format(subject, response.text))
        if current_app.config['ENVIRONMENT'] != 'live':  # pragma: no cover
            current_app.logger.info('Email to: {}'.format(to))
            current_app.logger.info('Email provider: {}'.format(email_provider.api_url))
            current_app.logger.info('Email key: {}'.format(email_provider.api_key[:5]))

        return response.status_code
    else:
        data = {
            'to': to,
            'from_email': from_email,
            'from_name': from_name,
            'subject': subject,
            'message': message
        }
        current_app.logger.info('No email providers configured, email would have sent: {}'.format(data))
<commit_msg>Add some logging to investigate email<commit_after>from flask import current_app, jsonify, render_template
import json
import os
import re
import requests
from html.parser import HTMLParser

from na_common.dates import get_nice_event_dates

from app.comms.encryption import encrypt
from app.errors import InvalidRequest
from app.models import BASIC, EVENT, MAGAZINE
from app.dao.events_dao import dao_get_event_by_id
from app.dao.emails_dao import dao_get_past_hour_email_count_for_provider, dao_get_todays_email_count_for_provider
from app.dao.email_providers_dao import dao_get_first_email_provider, dao_get_next_email_provider
from app.dao.magazines_dao import dao_get_magazine_by_id

h = HTMLParser()

def get_email_provider(override=False, email_provider=None):
    if not email_provider:
        email_provider = dao_get_first_email_provider()
        if not email_provider:
            return None

    daily_email_count = dao_get_todays_email_count_for_provider(email_provider.id)
    hourly_email_count = 0

    if email_provider.hourly_limit > 0:
        hourly_email_count = dao_get_past_hour_email_count_for_provider(email_provider.id)

    if daily_email_count > email_provider.daily_limit:
        if override:
            next_email_provider = dao_get_next_email_provider(email_provider.pos)

            if next_email_provider:
                return get_email_provider(override, next_email_provider)
        else:
            email_provider.limit = 0
            email_provider.daily_limit_reached = True
        return email_provider

    if hourly_email_count > email_provider.hourly_limit:
        if override:
            next_email_provider = dao_get_next_email_provider(email_provider.pos)
            if next_email_provider:
                return get_email_provider(override, next_email_provider)
        email_provider.limit = 0
        email_provider.hourly_limit_reached = True

    if email_provider.hourly_limit > 0:
        email_provider.limit = email_provider.hourly_limit - hourly_email_count

    email_provider.limit = email_provider.daily_limit - daily_email_count
    return email_provider

def get_email_html(email_type, **kwargs):
    if email_type == EVENT:
        event = dao_get_event_by_id(kwargs.get('event_id'))
        member_id = kwargs.get('member_id')
        if not member_id:
            member_id = '0'  # for preview of emails
        current_app.logger.debug('Email Tokens %s', current_app.config['EMAIL_TOKENS'])
        unsubcode = encrypt(
            "{}={}".format(current_app.config['EMAIL_TOKENS']['member_id'], member_id),
            current_app.config['EMAIL_UNSUB_SALT']
        )
        return render_template(
            'emails/events.html',
            event=event,
            event_dates=get_nice_event_dates(event.event_dates),
            description=h.unescape(event.description),
            details=kwargs.get('details'),
            extra_txt=kwargs.get('extra_txt'),
            unsubcode=unsubcode
        )
    elif email_type == MAGAZINE:
        magazine = dao_get_magazine_by_id(kwargs.get('magazine_id'))
        topics = []
        if magazine.topics:
            _topics = [(t.split(':')[0], t.split(':')[1]) for t in magazine.topics.split('\n')]

            for title, description in _topics:
                topics.append({'title': title, 'description': description})
        return render_template(
            'emails/magazine.html',
            magazine=magazine,
            topics=topics
        )
    elif email_type == BASIC:
        return render_template(
            'emails/basic.html',
            title=kwargs.get('title', ''),
            message=kwargs.get('message')
        )

def get_email_data(data_map, to, subject, message, from_email, from_name):
    data_struct = {}
    if isinstance(to, list):
        to = ','.join(to)

    def set_data(_map, val):
        elements = _map.split(",")

        if len(elements) == 2:
            match = re.match(r"\[(?P<list_element>.+)\]", elements[1])
            if match:
                data_struct[elements[0]] = []

                for item in val.split(','):
                    data_struct[elements[0]].append({
                        match.group('list_element'): item
                    })
            else:
                if data_struct.get(elements[0]):
                    data_struct[elements[0]][elements[1]] = val
                else:
                    data_struct[elements[0]] = {elements[1]: val}
        else:
            data_struct[elements[0]] = val

    set_data(data_map['to'], to)
    set_data(data_map['subject'], subject)
    set_data(data_map['from'], from_email)
    if 'from_name' in data_map.keys():
        set_data(data_map['from_name'], from_name)
    set_data(data_map['message'], message)

    return data_struct

def send_email(to, subject, message, from_email=None, from_name=None, override=False):
    if current_app.config['EMAIL_DISABLED']:
        current_app.logger.info("Emails disabled, unset EMAIL_DISABLED env var to re-enable")
        return 200
    if current_app.config['ENVIRONMENT'] != 'live' or current_app.config.get('EMAIL_RESTRICT'):
        message = message.replace('', '
Test email, intended for {}
'.format(to))
        to = current_app.config['TEST_EMAIL']

    if not from_email:
        from_email = 'noreply@{}'.format(current_app.config['EMAIL_DOMAIN'])
    if not from_name:
        from_name = 'New Acropolis'

    email_provider = get_email_provider(override)

    if email_provider:
        if hasattr(email_provider, "hourly_limit_reached"):
            raise InvalidRequest('Hourly limit reached', 429)
        elif hasattr(email_provider, "daily_limit_reached"):
            raise InvalidRequest('Daily limit reached', 429)

        # temp log to investigate emails
        current_app.logger.info(f"Sending email to: {to[:10]}")

        data = get_email_data(email_provider.data_map, to, subject, message, from_email, from_name)
        data = data if email_provider.as_json else json.dumps(data)

        headers = {
            'api-key': email_provider.api_key,
            'accept': 'application/json',
            'content-type': 'application/json'
        } if email_provider.headers else None

        response = requests.post(
            email_provider.api_url,
            auth=('api', email_provider.api_key),
            headers=headers,
            data=data,
        )

        response.raise_for_status()
        current_app.logger.info('Sent email: {}, response: {}'.format(subject, response.text))
        if current_app.config['ENVIRONMENT'] != 'live':  # pragma: no cover
            current_app.logger.info('Email to: {}'.format(to))
            current_app.logger.info('Email provider: {}'.format(email_provider.api_url))
            current_app.logger.info('Email key: {}'.format(email_provider.api_key[:5]))

        return response.status_code
    else:
        data = {
            'to': to,
            'from_email': from_email,
            'from_name': from_name,
            'subject': subject,
            'message': message
        }
        current_app.logger.info('No email providers configured, email would have sent: {}'.format(data))
<|endoftext|>


<commit_before>#!/usr/bin/env python2
import sys, os, json, errno, subprocess, yum

class yumoter:
    def __init__(self, configFile, repobasepath):
        self.urlbasepath = "http://yumoter.gnmedia.net"
        self.repobasepath = repobasepath
        self.reloadConfig(configFile)
        self.yb = yum.YumBase()
        self.yb.setCacheDir(yum.misc.getCacheDir())
        self.yb.repos.disableRepo("*")

    def _getConfig(self, jsonFile):
        fh = open(jsonFile, 'r')
        jsonOutput = json.load(fh)
        fh.close()
        return jsonOutput

    def _mkdir_p(self, path):
        try:
            os.makedirs(path)
        except OSError as exc:
            if exc.errno == errno.EEXIST and os.path.isdir(path):
                pass
            else:
                raise

    def _translateToMajorVer(self, ver):
        # I: 6.4
        # O: 6
        if '.' in ver:
            ver = ver.split('.')[0]
        return ver

    def _getPaths(self):
        for repo in self.repoConfig:
            repopath = []
            # Does this repo have a path for promotion?
            if 'promotionpath' in self.repoConfig[repo]:
                for promopath in self.repoConfig[repo]['promotionpath']:
                    repopath.append("%s/%s/%s" % (self.repobasepath, self.repoConfig[repo]['path'], promopath))
            else:
                # repo does not have a path for promotion
                repopath.append("%s/%s" % (self.repobasepath, self.repoConfig[repo]['path']))
            self.repoConfig[repo]['fullpaths'] = repopath

    def _getURLs(self):
        for repo in self.repoConfig:
            repoURLs = []
            if 'promotionpath' in self.repoConfig[repo]:
                for promopath in self.repoConfig[repo]['promotionpath']:
                    repoURLs.append("%s/%s/%s" % (self.urlbasepath, self.repoConfig[repo]['path'], promopath))
            else:
                repoURLs.append("%s/%s" % (self.urlbasepath, self.repoConfig[repo]['path']))
            self.repoConfig[repo]['fullurls'] = repoURLs

    def _mkPaths(self):
        masterPathList = []
        for repo in self.repoConfig:
            if 'fullpaths' in self.repoConfig[repo]:
                for entry in self.repoConfig[repo]['fullpaths']:
                    masterPathList.append(entry)
        for entry in masterPathList:
            if not os.path.isdir(entry):
                print "creating missing dir: %s" % entry
                self._mkdir_p(entry)

    def _runRsync(self, rsrc, rdst, args):
        # str(rsrc), str(rdst), list(args)
        sysCall = ['rsync'] + args + [rsrc, rdst]
        print 'sysCall', sysCall
        rsyncStdout = []
        rsyncStderr = []
        p = subprocess.Popen(sysCall, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
        for line in iter(p.stdout.readline, ""):
            stdoutLine = line.strip() + '\r\n'
            rsyncStdout.append(stdoutLine)
            sys.stdout.write(stdoutLine)
            sys.stdout.flush()
        for line in iter(p.stderr.readline, ""):
            stderrLine = line.strip() + '\r\n'
            rsyncStderr.append(stderrLine)
            sys.stderr.write(stderrLine)
            sys.stderr.flush()
        return (rsyncStdout, rsyncStderr)
        # TODO check return status please. Stop coding like a 12 year old.

    def _loadRepo(self, reponame, repo):
        if reponame == "os-64":
            print "Adding repo:", reponame, repo
            self.yb.add_enable_repo(reponame, baseurls=[repo])

    def _returnNewestByNameArch(self, patternsList):
        pkgs = self.yb.pkgSack.returnNewestByNameArch(patterns=patternsList)
        for pkg in pkgs:
            print "%s: %s" % (pkg, pkg.summary)

    def loadRepos(self, osVer, env):
        # this should load all the repos for osVer in env
        # Should use an internal method to load one repo
        # by url.
        loadrepos = []
        shortenv = self._translateToMajorVer(env)
        for repo in self.repoConfig:
            if '.' in self.repoConfig[repo]['osver']:
                # This is a repo sensative to minor versions
                if osVer == self.repoConfig[repo]['osver']:
                    loadrepos.append(repo)
            else:
                # This repo cares only about major versions
                if self._translateToMajorVer(osVer) == self.repoConfig[repo]['osver']:
                    loadrepos.append(repo)
        for repo in loadrepos:
            if len(self.repoConfig[repo]['fullurls']) == 1:
                # This repo only has one env. Repo is not promoted.
                # Load the only URL you can
                self._loadRepo(repo, self.repoConfig[repo]['fullurls'][0])
            else:
                self._loadRepo(repo, self.repoConfig[repo]['fullurls'][self.repoConfig[repo]['promotionpath'].index(env)])

    def getDeps(self, pkgObj):
        depDicts = yb.findDeps([pkgObj])
        return depDicts

    def syncRepos(self):
        outputList = []
        for repo in self.repoConfig:
            # Only repos with upstream set need to be synced.
            if 'upstream' in self.repoConfig[repo]:
                # If the dst dir doesn't exist, create it.
                if not os.path.isdir(self.repoConfig[repo]['fullpaths'][0]):
                    self._mkPaths()
                outputList.append(self._runRsync(self.repoConfig[repo]['upstream'], self.repoConfig[repo]['fullpaths'][0], ['-av', '--progress']))
        return outputList

    def reloadConfig(self, jsonFile):
        self.repoConfig = self._getConfig(jsonFile)
        self._getPaths()
        self._getURLs()

'''
    def repoSearch(self, pkgName, repos):

        pkgs = yb.pkgSack.returnNewestByNameArch(patterns=pkgName)
        for pkg in pkgs:
            print "%s: %s" % (pkg, pkg.summary)
'''<commit_msg>trying again<commit_after>#!/usr/bin/env python2
import sys, os, json, errno, subprocess, yum

class yumoter:
    def __init__(self, configFile, repobasepath):
        self.urlbasepath = "http://yumoter.gnmedia.net"
        self.repobasepath = repobasepath
        self.reloadConfig(configFile)
        self.yb = yum.YumBase()
        self.yb.setCacheDir(yum.misc.getCacheDir())
        self.yb.repos.disableRepo("*")

    def _getConfig(self, jsonFile):
        fh = open(jsonFile, 'r')
        jsonOutput = json.load(fh)
        fh.close()
        return jsonOutput

    def _mkdir_p(self, path):
        try:
            os.makedirs(path)
        except OSError as exc:
            if exc.errno == errno.EEXIST and os.path.isdir(path):
                pass
            else:
                raise

    def _translateToMajorVer(self, ver):
        # I: 6.4
        # O: 6
        if '.' in ver:
            ver = ver.split('.')[0]
        return ver

    def _getPaths(self):
        for repo in self.repoConfig:
            repopath = []
            # Does this repo have a path for promotion?
            if 'promotionpath' in self.repoConfig[repo]:
                for promopath in self.repoConfig[repo]['promotionpath']:
                    repopath.append("%s/%s/%s" % (self.repobasepath, self.repoConfig[repo]['path'], promopath))
            else:
                # repo does not have a path for promotion
                repopath.append("%s/%s" % (self.repobasepath, self.repoConfig[repo]['path']))
            self.repoConfig[repo]['fullpaths'] = repopath

    def _getURLs(self):
        for repo in self.repoConfig:
            repoURLs = []
            if 'promotionpath' in self.repoConfig[repo]:
                for promopath in self.repoConfig[repo]['promotionpath']:
                    repoURLs.append("%s/%s/%s" % (self.urlbasepath, self.repoConfig[repo]['path'], promopath))
            else:
                repoURLs.append("%s/%s" % (self.urlbasepath, self.repoConfig[repo]['path']))
            self.repoConfig[repo]['fullurls'] = repoURLs

    def _mkPaths(self):
        masterPathList = []
        for repo in self.repoConfig:
            if 'fullpaths' in self.repoConfig[repo]:
                for entry in self.repoConfig[repo]['fullpaths']:
                    masterPathList.append(entry)
        for entry in masterPathList:
            if not os.path.isdir(entry):
                print "creating missing dir: %s" % entry
                self._mkdir_p(entry)

    def _runRsync(self, rsrc, rdst, args):
        # str(rsrc), str(rdst), list(args)
        sysCall = ['rsync'] + args + [rsrc, rdst]
        print 'sysCall', sysCall
        rsyncStdout = []
        rsyncStderr = []
        p = subprocess.Popen(sysCall, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
        for line in iter(p.stdout.readline, ""):
            stdoutLine = line.strip() + '\r\n'
            rsyncStdout.append(stdoutLine)
            sys.stdout.write(stdoutLine)
            sys.stdout.flush()
        for line in iter(p.stderr.readline, ""):
            stderrLine = line.strip() + '\r\n'
            rsyncStderr.append(stderrLine)
            sys.stderr.write(stderrLine)
            sys.stderr.flush()
        return (rsyncStdout, rsyncStderr)
        # TODO check return status please. Stop coding like a 12 year old.

    def _loadRepo(self, reponame, repo):
        if reponame == "os-64":
            print "Adding repo:", reponame, repo
            self.yb.add_enable_repo(reponame, baseurls=[repo])

    def _returnNewestByNameArch(self, patternsList):
        print patternsList
        pkgs = self.yb.pkgSack.returnNewestByNameArch(patterns=patternsList)
        for pkg in pkgs:
            print "%s: %s" % (pkg, pkg.summary)

    def loadRepos(self, osVer, env):
        # this should load all the repos for osVer in env
        # Should use an internal method to load one repo
        # by url.
        loadrepos = []
        shortenv = self._translateToMajorVer(env)
        for repo in self.repoConfig:
            if '.' in self.repoConfig[repo]['osver']:
                # This is a repo sensative to minor versions
                if osVer == self.repoConfig[repo]['osver']:
                    loadrepos.append(repo)
            else:
                # This repo cares only about major versions
                if self._translateToMajorVer(osVer) == self.repoConfig[repo]['osver']:
                    loadrepos.append(repo)
        for repo in loadrepos:
            if len(self.repoConfig[repo]['fullurls']) == 1:
                # This repo only has one env. Repo is not promoted.
                # Load the only URL you can
                self._loadRepo(repo, self.repoConfig[repo]['fullurls'][0])
            else:
                self._loadRepo(repo, self.repoConfig[repo]['fullurls'][self.repoConfig[repo]['promotionpath'].index(env)])

    def getDeps(self, pkgObj):
        depDicts = yb.findDeps([pkgObj])
        return depDicts

    def syncRepos(self):
        outputList = []
        for repo in self.repoConfig:
            # Only repos with upstream set need to be synced.
            if 'upstream' in self.repoConfig[repo]:
                # If the dst dir doesn't exist, create it.
                if not os.path.isdir(self.repoConfig[repo]['fullpaths'][0]):
                    self._mkPaths()
                outputList.append(self._runRsync(self.repoConfig[repo]['upstream'], self.repoConfig[repo]['fullpaths'][0], ['-av', '--progress']))
        return outputList

    def reloadConfig(self, jsonFile):
        self.repoConfig = self._getConfig(jsonFile)
        self._getPaths()
        self._getURLs()

'''
    def repoSearch(self, pkgName, repos):

        pkgs = yb.pkgSack.returnNewestByNameArch(patterns=pkgName)
        for pkg in pkgs:
            print "%s: %s" % (pkg, pkg.summary)
'''<|endoftext|>


<commit_before>import networkx as nx
import numpy as np
import os.path
import math
import pickle
from sklearn import linear_model

import config
import util
from balanced_partition import BalancedPartition
import link_rank_clustering as lrc

def _remove_unidirectional_edges(G):
	adjmat = nx.adjacency_matrix(G)
	adjmat = adjmat + adjmat.T
	adjmat[adjmat == 1] = 0
	adjmat[adjmat == 2] = 1
	G = nx.from_scipy_sparse_matrix(adjmat, create_using=nx.Graph())
	return G

def _get_true_ate(G, adjmat):
	N = adjmat.shape[0]
	Z0 = np.zeros(N)
	Z1 = np.ones(N)
	Y0 = util.outcome_generator(G, adjmat, Z0)
	Y1 = util.outcome_generator(G, adjmat, Z1)
	return np.mean(Y1 - Y0)

def _sampling(G, model, method):
	N = G.number_of_nodes()
	if method == "uniform":
		return np.random.binomial(1, 0.5, N)
	mapping = {"b1": 1, "b2": 2, "LRC": 3}
	cluster_type = mapping[method]
	Z = np.empty(N)
	name = config.dynamic["graph_name"]
	cluster_cache_path = util.get_file_path("cluster_cache", name=name, cluster_type=cluster_type)
	if os.path.exists(cluster_cache_path):
		print("Loading clusters from cache...")
		clusters = pickle.load(open(cluster_cache_path, "rb"))
	else:
		if cluster_type == mapping["b2"]:
			BP = BalancedPartition(G)
			clusters = BP.clustering()
		elif cluster_type == mapping["LRC"]:
			clusters = lrc.clustering(G)
		util.save_cluster(clusters, name, cluster_type)
	cnt = 0
	for cluster in clusters:
		assignment = np.random.binomial(1, 0.5)
		cluster = list(cluster)
		Z[cluster] = assignment
		cnt += len(cluster)
	assert cnt == N, "cnt not equals N"
	return Z

def _estimate_lm1(Z, sigma, outcome):
	assert type(Z) == np.ndarray and Z.ndim == 1
	assert type(sigma) == np.ndarray and sigma.ndim == 1
	assert type(outcome) == np.ndarray and outcome.ndim == 1
	lr = linear_model.LinearRegression()
	X = np.stack((Z, sigma), axis=-1)
	reg = lr.fit(X, outcome)
	beta = reg.coef_[0]
	gamma = reg.coef_[1]
	return beta + gamma

def _estimate_lm2(Z, sigma, outcome):
	assert type(Z) == np.ndarray and Z.ndim == 1
	assert type(sigma) == np.ndarray and sigma.ndim == 1
	assert type(outcome) == np.ndarray and outcome.ndim == 1
	lr = linear_model.LinearRegression()
	sigma0, outcome0 = sigma[Z == 0].reshape(-1, 1), outcome[Z == 0].reshape(-1, 1)
	sigma1, outcome1 = sigma[Z == 1].reshape(-1, 1), outcome[Z == 1].reshape(-1, 1)
	reg0 = lr.fit(sigma0, outcome0)
	reg1 = lr.fit(sigma1, outcome1)
	alpha0 = reg0.intercept_
	alpha1 = reg1.intercept_
	gamma1 = reg1.coef_[0]
	return alpha1 + gamma1 - alpha0

# def _convert(G, method):
# 	if method == "baseline1":
# 		G = G.to_undirected()
# 	adjmat = nx.adjacency_matrix(G)
# 	return G, adjmat

def estimate(G, model, method):
	assert G.__class__.__name__ == "DiGraph", "Graph isn't digraph"
	assert method in ["uniform", "b1", "LRC"], "Method provided (%s) not exists" % method

	#G, adjmat = _convert(G, method)
	adjmat = nx.adjacency_matrix(G)
	Z = _sampling(G, model, method)
	outcome = util.outcome_generator(G, adjmat, Z)
	true_ate = _get_true_ate(G, adjmat)
	if method == "uniform":
		estimated_ate = np.mean(outcome[Z==1]) - np.mean(outcome[Z==0])
	elif model == "linear1":
		sigma = util.treated_proportion(adjmat, Z)
		estimated_ate = _estimate_lm1(Z, sigma, outcome)
	elif model == "linear2":
		sigma = util.treated_proportion(adjmat, Z)
		estimated_ate = _estimate_lm2(Z, sigma, outcome)
	else:
		raise Exception("Model specified not exists")
	return true_ate, estimated_ate
<commit_msg>bug fix<commit_after>import networkx as nx
import numpy as np
import os.path
import math
import pickle
from sklearn import linear_model

import config
import util
from balanced_partition import BalancedPartition
import link_rank_clustering as lrc

def _remove_unidirectional_edges(G):
	adjmat = nx.adjacency_matrix(G)
	adjmat = adjmat + adjmat.T
	adjmat[adjmat == 1] = 0
	adjmat[adjmat == 2] = 1
	G = nx.from_scipy_sparse_matrix(adjmat, create_using=nx.Graph())
	return G

def _get_true_ate(G, adjmat):
	N = adjmat.shape[0]
	Z0 = np.zeros(N)
	Z1 = np.ones(N)
	Y0 = util.outcome_generator(G, adjmat, Z0)
	Y1 = util.outcome_generator(G, adjmat, Z1)
	return np.mean(Y1 - Y0)

def _sampling(G, model, method):
	N = G.number_of_nodes()
	if method == "uniform":
		return np.random.binomial(1, 0.5, N)
	mapping = {"b1": 1, "LRC": 1}
	cluster_type = mapping[method]
	Z = np.empty(N)
	name = config.dynamic["graph_name"]
	cluster_cache_path = util.get_file_path("cluster_cache", name=name, cluster_type=cluster_type)
	if os.path.exists(cluster_cache_path):
		print("Loading clusters from cache...")
		clusters = pickle.load(open(cluster_cache_path, "rb"))
	else:
		if cluster_type == mapping["b1"]:
			BP = BalancedPartition(G)
			clusters = BP.clustering()
		elif cluster_type == mapping["LRC"]:
			clusters = lrc.clustering(G)
		util.save_cluster(clusters, name, cluster_type)
	cnt = 0
	for cluster in clusters:
		assignment = np.random.binomial(1, 0.5)
		cluster = list(cluster)
		Z[cluster] = assignment
		cnt += len(cluster)
	assert cnt == N, "cnt not equals N"
	return Z

def _estimate_lm1(Z, sigma, outcome):
	assert type(Z) == np.ndarray and Z.ndim == 1
	assert type(sigma) == np.ndarray and sigma.ndim == 1
	assert type(outcome) == np.ndarray and outcome.ndim == 1
	lr = linear_model.LinearRegression()
	X = np.stack((Z, sigma), axis=-1)
	reg = lr.fit(X, outcome)
	beta = reg.coef_[0]
	gamma = reg.coef_[1]
	return beta + gamma

def _estimate_lm2(Z, sigma, outcome):
	assert type(Z) == np.ndarray and Z.ndim == 1
	assert type(sigma) == np.ndarray and sigma.ndim == 1
	assert type(outcome) == np.ndarray and outcome.ndim == 1
	lr = linear_model.LinearRegression()
	sigma0, outcome0 = sigma[Z == 0].reshape(-1, 1), outcome[Z == 0].reshape(-1, 1)
	sigma1, outcome1 = sigma[Z == 1].reshape(-1, 1), outcome[Z == 1].reshape(-1, 1)
	reg0 = lr.fit(sigma0, outcome0)
	reg1 = lr.fit(sigma1, outcome1)
	alpha0 = reg0.intercept_
	alpha1 = reg1.intercept_
	gamma1 = reg1.coef_[0]
	return alpha1 + gamma1 - alpha0

# def _convert(G, method):
# 	if method == "baseline1":
# 		G = G.to_undirected()
# 	adjmat = nx.adjacency_matrix(G)
# 	return G, adjmat

def estimate(G, model, method):
	assert G.__class__.__name__ == "DiGraph", "Graph isn't digraph"
	assert method in ["uniform", "b1", "LRC"], "Method provided (%s) not exists" % method

	#G, adjmat = _convert(G, method)
	adjmat = nx.adjacency_matrix(G)
	Z = _sampling(G, model, method)
	outcome = util.outcome_generator(G, adjmat, Z)
	true_ate = _get_true_ate(G, adjmat)
	if method == "uniform":
		estimated_ate = np.mean(outcome[Z==1]) - np.mean(outcome[Z==0])
	elif model == "linear1":
		sigma = util.treated_proportion(adjmat, Z)
		estimated_ate = _estimate_lm1(Z, sigma, outcome)
	elif model == "linear2":
		sigma = util.treated_proportion(adjmat, Z)
		estimated_ate = _estimate_lm2(Z, sigma, outcome)
	else:
		raise Exception("Model specified not exists")
	return true_ate, estimated_ate
<|endoftext|>


<commit_before>import os
import djcelery

# Debug settings
DEBUG = False
TEMPLATE_DEBUG = True

# Zone Settings
TIME_ZONE = 'UTC'
LANGUAGE_CODE = 'en-us'
SITE_ID = 1
USE_I18N = True

# Defines the Static Media storage as per staticfiles contrib
STATIC_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', 'static')
STATIC_URL = '/static/'
ADMIN_MEDIA_PREFIX = STATIC_URL + 'admin/'

# Make this unique, and don't share it with anybody.
SECRET_KEY = ''

# List of callables that know how to import templates from various sources.
TEMPLATE_LOADERS = (
    'django.template.loaders.filesystem.Loader',
    'django.template.loaders.app_directories.Loader',
)

MIDDLEWARE_CLASSES = (
    'django.middleware.common.CommonMiddleware',
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'sso.middleware.InactiveLogoutMiddleware',
    'pagination.middleware.PaginationMiddleware',
    'sso.middleware.IGBMiddleware',
    'sso.middleware.IPTrackingMiddleware',
)

ROOT_URLCONF = 'urls'

TEMPLATE_DIRS = (
    os.path.join(os.path.dirname(__file__), '..', 'templates'),
)

TEMPLATE_CONTEXT_PROCESSORS = (
    "django.contrib.auth.context_processors.auth",
    "django.core.context_processors.debug",
    "django.core.context_processors.i18n",
    "django.core.context_processors.static",
    "django.contrib.messages.context_processors.messages",
    "django.core.context_processors.request"
)

MESSAGE_STORAGE = 'django.contrib.messages.storage.session.SessionStorage'

INSTALLED_APPS = (
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.messages',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.sites',
    'django.contrib.humanize',
    'django.contrib.staticfiles',
    'nexus',
    'gargoyle',
    'sentry',
    'sentry.client',
    'south',
    'piston',
    'djcelery',
    'registration',
    'pagination',
    'eve_proxy',
    'eve_api',
    'reddit',
    'hr',
    'sso',
    'groups',
    'api',
    'tools',
)

AUTHENTICATION_BACKENDS = (
    'sso.backends.SimpleHashModelBackend',
)

AUTH_PROFILE_MODULE = 'sso.SSOUser'
LOGIN_REDIRECT_URL = "/profile"
LOGIN_URL = "/login"

### OAuth

OAUTH_AUTH_VIEW = 'api.views.oauth_auth_view'
OAUTH_CALLBACK_VIEW = 'api.views.oauth_callback_view'

### Celery Schedule

from celeryschedule import *

CELERY_SEND_TASK_ERROR_EMAILS = True
CELERY_RESULT_BACKEND = "amqp"
CELERY_DISABLE_RATE_LIMITS = True
CELERYD_PREFETCH_MULTIPLIER = 128
CELERY_ALWAYS_EAGER = DEBUG
CELERY_EAGER_PROPAGATES_EXCEPTIONS = DEBUG

# Load the Celery tasks
djcelery.setup_loader()

GARGOYLE_AUTO_CREATE = True

# Switches for Gargoyle
GARGOYLE_SWITCH_DEFAULTS = {
    'reddit': {
      'is_active': True,
      'label': 'Reddit Functionality',
      'description': 'Enables/Disables the Reddit integration for HR and SSO.',
    },
    'hr': {
      'is_active': True,
      'label': 'HR Functions',
      'description': 'Enables/Disables the HR functionality.',
    },
    'eve-cak': {
      'is_active': False,
      'label': 'EVE Customizable API Keys',
      'description': 'Enables/Disables EVE API CAK support.',
    },
    'eve-testapi': {
      'is_active': False,
      'label': 'EVE Test API Endpoints',
      'description': 'Use the Test API endpoints instead of Live.',
    },
    'eve-keydelete': {
      'is_active': False,
      'label': 'Allow EVE API Key Delete',
      'description': 'Allows API keys to be deleted by th end user.',
    },
    'api-disableprocessing': {
      'is_active': False,
      'label': 'Disable API Backend Processing',
      'description': 'Disables backend processing for the EVE API, stops Auth hammering the API during outages',
    }

}

LOGGING = {
    'version': 1,
    'disable_existing_loggers': True,
    'handlers': {
        'null': {
            'level':'DEBUG',
            'class':'django.utils.log.NullHandler',
        },
        'mail_admins': {
            'level': 'ERROR',
            'class': 'django.utils.log.AdminEmailHandler',
        },
        'sentry': {
            'level': 'DEBUG',
            'class': 'sentry.client.handlers.SentryHandler',
        },
    },
    'loggers': {
        '()': {
            'level': 'WARNING',
            'handlers': ['sentry'],
        },
        'sentry.errors': {
            'level': 'DEBUG',
            'handlers': ['null'],
            'propagate': False,
        },
        'django': {
            'handlers':['null'],
            'propagate': True,
            'level':'INFO',
        },
        'django.request': {
            'handlers': ['mail_admins'],
            'level': 'ERROR',
            'propagate': False,
        },
    }
}
<commit_msg>Propagate django.request errors down to sentry instead of mail, also fix global logging entry<commit_after>import os
import djcelery

# Debug settings
DEBUG = False
TEMPLATE_DEBUG = True

# Zone Settings
TIME_ZONE = 'UTC'
LANGUAGE_CODE = 'en-us'
SITE_ID = 1
USE_I18N = True

# Defines the Static Media storage as per staticfiles contrib
STATIC_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', 'static')
STATIC_URL = '/static/'
ADMIN_MEDIA_PREFIX = STATIC_URL + 'admin/'

# Make this unique, and don't share it with anybody.
SECRET_KEY = ''

# List of callables that know how to import templates from various sources.
TEMPLATE_LOADERS = (
    'django.template.loaders.filesystem.Loader',
    'django.template.loaders.app_directories.Loader',
)

MIDDLEWARE_CLASSES = (
    'django.middleware.common.CommonMiddleware',
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'sso.middleware.InactiveLogoutMiddleware',
    'pagination.middleware.PaginationMiddleware',
    'sso.middleware.IGBMiddleware',
    'sso.middleware.IPTrackingMiddleware',
)

ROOT_URLCONF = 'urls'

TEMPLATE_DIRS = (
    os.path.join(os.path.dirname(__file__), '..', 'templates'),
)

TEMPLATE_CONTEXT_PROCESSORS = (
    "django.contrib.auth.context_processors.auth",
    "django.core.context_processors.debug",
    "django.core.context_processors.i18n",
    "django.core.context_processors.static",
    "django.contrib.messages.context_processors.messages",
    "django.core.context_processors.request"
)

MESSAGE_STORAGE = 'django.contrib.messages.storage.session.SessionStorage'

INSTALLED_APPS = (
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.messages',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.sites',
    'django.contrib.humanize',
    'django.contrib.staticfiles',
    'nexus',
    'gargoyle',
    'sentry',
    'sentry.client',
    'south',
    'piston',
    'djcelery',
    'registration',
    'pagination',
    'eve_proxy',
    'eve_api',
    'reddit',
    'hr',
    'sso',
    'groups',
    'api',
    'tools',
)

AUTHENTICATION_BACKENDS = (
    'sso.backends.SimpleHashModelBackend',
)

AUTH_PROFILE_MODULE = 'sso.SSOUser'
LOGIN_REDIRECT_URL = "/profile"
LOGIN_URL = "/login"

### OAuth

OAUTH_AUTH_VIEW = 'api.views.oauth_auth_view'
OAUTH_CALLBACK_VIEW = 'api.views.oauth_callback_view'

### Celery Schedule

from celeryschedule import *

CELERY_SEND_TASK_ERROR_EMAILS = True
CELERY_RESULT_BACKEND = "amqp"
CELERY_DISABLE_RATE_LIMITS = True
CELERYD_PREFETCH_MULTIPLIER = 128
CELERY_ALWAYS_EAGER = DEBUG
CELERY_EAGER_PROPAGATES_EXCEPTIONS = DEBUG

# Load the Celery tasks
djcelery.setup_loader()

GARGOYLE_AUTO_CREATE = True

# Switches for Gargoyle
GARGOYLE_SWITCH_DEFAULTS = {
    'reddit': {
      'is_active': True,
      'label': 'Reddit Functionality',
      'description': 'Enables/Disables the Reddit integration for HR and SSO.',
    },
    'hr': {
      'is_active': True,
      'label': 'HR Functions',
      'description': 'Enables/Disables the HR functionality.',
    },
    'eve-cak': {
      'is_active': False,
      'label': 'EVE Customizable API Keys',
      'description': 'Enables/Disables EVE API CAK support.',
    },
    'eve-testapi': {
      'is_active': False,
      'label': 'EVE Test API Endpoints',
      'description': 'Use the Test API endpoints instead of Live.',
    },
    'eve-keydelete': {
      'is_active': False,
      'label': 'Allow EVE API Key Delete',
      'description': 'Allows API keys to be deleted by th end user.',
    },
    'api-disableprocessing': {
      'is_active': False,
      'label': 'Disable API Backend Processing',
      'description': 'Disables backend processing for the EVE API, stops Auth hammering the API during outages',
    }

}

LOGGING = {
    'version': 1,
    'disable_existing_loggers': True,
    'handlers': {
        'null': {
            'level':'DEBUG',
            'class':'django.utils.log.NullHandler',
        },
        'mail_admins': {
            'level': 'ERROR',
            'class': 'django.utils.log.AdminEmailHandler',
        },
        'sentry': {
            'level': 'DEBUG',
            'class': 'sentry.client.handlers.SentryHandler',
        },
    },
    'loggers': {
        '': {
            'level': 'WARNING',
            'handlers': ['sentry'],
        },
        'sentry.errors': {
            'level': 'DEBUG',
            'handlers': ['null'],
            'propagate': False,
        },
        'django': {
            'handlers':['null'],
            'propagate': True,
            'level':'INFO',
        },
        'django.request': {
            'handlers': ['null'],
            'level': 'ERROR',
            'propagate': True,
        },
    }
}
<|endoftext|>


<commit_before>import logging
import os.path
import shutil
from datetime import datetime
from pathlib import Path
from typing import Dict, Callable

from . import Checks
from . import Utils
from .GoogleAlbumMedia import GoogleAlbumMedia
from .GoogleAlbumsRow import GoogleAlbumsRow
from .GooglePhotosMedia import GooglePhotosMedia
from .GooglePhotosRow import GooglePhotosRow
from .LocalData import LocalData
from .restclient import RestClient

log = logging.getLogger(__name__)

PAGE_SIZE = 100
ALBUM_ITEMS = 50

class GoogleAlbumsSync(object):
    """A Class for managing the indexing and download Google of Albums
    """

    def __init__(self, api: RestClient, root_folder: Path, db: LocalData,
                 flush: bool, photos_path: Path,
                 albums_path: Path, use_flat_path=True, omit_album_date=True,
                 use_hardlinks=False):
        """
        Parameters:
            root_folder: path to the root of local file synchronization
            api: object representing the Google REST API
            db: local database for indexing
            :param use_flat_path:
            :param omit_album_date:
            :param use_hardlinks:
            :param photos_path:
            :param albums_path:
        """
        self._root_folder: Path = root_folder
        self._photos_folder = Path(photos_path)
        self._albums_folder = Path(albums_path)
        self._links_root = self._root_folder / self._albums_folder
        self._photos_root = self._root_folder / self._photos_folder
        self._db: LocalData = db
        self._api: RestClient = api
        self._use_flat_path = use_flat_path
        self._omit_album_date = omit_album_date
        self._use_hardlinks = use_hardlinks
        self.flush = flush
        # these properties are set after construction
        self.album = None
        self.shared_albums = True
        self.album_index = True
        self.use_start_date = False
        self.favourites = False

    @classmethod
    def make_search_parameters(cls, album_id: str,
                               page_token: str = None) -> Dict:
        body = {
            'pageToken': page_token,
            'albumId': album_id,
            'pageSize': PAGE_SIZE
        }
        return body

    def fetch_album_contents(self, album_id: str,
                             add_media_items: bool) -> (datetime, datetime):
        first_date = Utils.maximum_date()
        last_date = Utils.MINIMUM_DATE
        body = self.make_search_parameters(album_id=album_id)
        response = self._api.mediaItems.search.execute(body)
        position = -1
        while response:
            items_json = response.json()
            media_json = items_json.get('mediaItems')
            # cope with empty albums
            if not media_json:
                if not items_json.get('nextPageToken'):
                    break
                else:
                    media_json = []
                    log.warning('*** Empty Media JSON with a Next Page Token')

            for media_item_json in media_json:
                position += 1
                media_item = GooglePhotosMedia(media_item_json)
                log.debug('----%s', media_item.filename)
                self._db.put_album_file(album_id, media_item.id, position)
                last_date = max(media_item.create_date, last_date)
                first_date = min(media_item.create_date, first_date)

                # this adds other users photos from shared albums
                # Todo - This will cause two copies of a file to appear for
                #  those shared items you have imported into your own library.
                #  They will have different RemoteIds, one will point to your
                #  library copy (you own) and one to the shared item in the
                #  the folder. Currently with the meta data available it would
                #  be impossible to eliminate these without eliminating other
                #  cases where date and filename (TITLE) match
                if add_media_items:
                    media_item.set_path_by_date(self._photos_folder,
                                                self._use_flat_path)
                    (num, row) = self._db.file_duplicate_no(
                        str(media_item.filename),
                        str(media_item.relative_folder),
                        media_item.id)
                    # we just learned if there were any duplicates in the db
                    media_item.duplicate_number = num

                    log.debug('Adding album media item %s %s %s',
                            media_item.relative_path, media_item.filename,
                            media_item.duplicate_number)
                    self._db.put_row(
                        GooglePhotosRow.from_media(media_item), False)

            next_page = items_json.get('nextPageToken')
            if next_page:
                body = self.make_search_parameters(album_id=album_id,
                                                   page_token=next_page)
                response = self._api.mediaItems.search.execute(body)
            else:
                break
        return first_date, last_date

    def index_album_media(self):
        # we now index all contents of non-shared albums due to the behaviour
        # reported here https://github.com/gilesknap/gphotos-sync/issues/89
        if self.shared_albums:
            self.index_albums_type(self._api.sharedAlbums.list.execute,
                                   'sharedAlbums', "Shared (titled) Albums",
                                   False, True)
        self.index_albums_type(self._api.albums.list.execute,
                               'albums', "Albums", True, self.album_index)

    def index_albums_type(self, api_function: Callable, item_key: str,
                          description: str, allow_null_title: bool,
                          add_media_items: bool):
        """
        query google photos interface for a list of all albums and index their
        contents into the db
        """
        log.warning('Indexing {} ...'.format(description))

        # when only looking for favourites do not download album contents
        if self.favourites:
            add_media_items = False

        # there are no filters in album listing at present so it always a
        # full rescan - it's quite quick
        count = 0
        response = api_function(pageSize=ALBUM_ITEMS)
        while response:
            results = response.json()
            for album_json in results.get(item_key, []):
                count += 1

                album = GoogleAlbumMedia(album_json)
                indexed_album = self._db.get_album(album_id=album.id)
                already_indexed = indexed_album.size == album.size if \
                    indexed_album else False

                if self.album and self.album != album.orig_name:
                    log.debug('Skipping Album: %s, photos: %d '
                              '(does not match --album)', album.filename,
                              album.size)
                elif not allow_null_title and album.description == 'none':
                    log.debug('Skipping no-title album, photos: %d',
                              album.size)
                elif already_indexed and not self.flush:
                    log.debug('Skipping Album: %s, photos: %d', album.filename,
                              album.size)
                else:
                    log.info('Indexing Album: %s, photos: %d', album.filename,
                             album.size)
                    first_date, last_date = self.fetch_album_contents(
                        album.id, add_media_items)
                    # write the album data down now we know the contents'
                    # date range
                    gar = GoogleAlbumsRow.from_parm(
                        album.id, album.filename, album.size,
                        first_date, last_date)
                    self._db.put_row(gar, update=indexed_album)

            next_page = results.get('nextPageToken')
            if next_page:
                response = api_function(pageSize=ALBUM_ITEMS,
                                        pageToken=next_page)
            else:
                break
        log.warning('Indexed %d %s', count, description)

    def album_folder_name(
        self, album_name: str, start_date: datetime, end_date: datetime
    ) -> Path:
        if self._omit_album_date:
            rel_path = album_name
        else:
            if self.use_start_date:
                d = start_date
            else:
                d = end_date
            year = Utils.safe_str_time(d, '%Y')
            month = Utils.safe_str_time(d, '%m%d')

            if self._use_flat_path:
                rel_path = u"{0}-{1} {2}".format(year, month, album_name)
            else:
                rel_path = Path(year) / u"{0} {1}".format(month, album_name)
        link_folder: Path = self._links_root / rel_path
        return link_folder

    def create_album_content_links(self):
        log.warning("Creating album folder links to media ...")
        count = 0
        album_item = 0
        current_rid = ''

        # always re-create all album links - it is quite fast and a good way
        # to ensure consistency
        # especially now that we have --album-date-by-first-photo
        if self._links_root.exists():
            log.debug('removing previous album links tree')
            shutil.rmtree(self._links_root)
        re_download = not self._links_root.exists()

        for (
                path, file_name, album_name,
                start_date_str, end_date_str,
                rid, created
        ) in \
                self._db.get_album_files(download_again=re_download):
            if current_rid == rid:
                album_item += 1
            else:
                self._db.put_album_downloaded(rid)
                current_rid = rid
                album_item = 0
            end_date = Utils.string_to_date(end_date_str)
            start_date = Utils.string_to_date(start_date_str)

            if len(str(self._root_folder / path)) > Checks.MAX_PATH_LENGTH:
                max_path_len = Checks.MAX_PATH_LENGTH - len(
                    str(self._root_folder)
                )
                log.debug('This path needs to be shrunk: %s' %
                          Path(self._root_folder / path))
                path = path[:max_path_len]
                log.debug('Shrunk to: %s' %
                          Path(self._root_folder / path))

            file_name = file_name[:Checks.MAX_FILENAME_LENGTH]

            full_file_name = self._root_folder / path / file_name

            link_folder: Path = self.album_folder_name(
                album_name, start_date, end_date
            )

            link_file = link_folder / "{:04d}_{}".format(album_item, file_name)
            # incredibly, pathlib.Path.relative_to cannot handle
            # '../' in a relative path !!! reverting to os.path
            relative_filename = os.path.relpath(full_file_name,
                                                str(link_folder))
            log.debug('adding album link %s -> %s', relative_filename,
                      link_file)
            try:
                if not link_folder.is_dir():
                    log.debug('new album folder %s', link_folder)
                    link_folder.mkdir(parents=True)

                created_date = Utils.string_to_date(created)
                if self._use_hardlinks:
                    if full_file_name.exists():
                        os.link(full_file_name, link_file)
                    else:
                        log.debug('skip hardlink for %s, not downloaded',
                                  file_name)
                else:
                    link_file.symlink_to(relative_filename)

                if link_file.exists():
                    count += 1
                    # Windows tries to follow symlinks even though we specify
                    # follow_symlinks=False. So disable setting of link date
                    # if follow not supported
                    if os.utime in os.supports_follow_symlinks:
                        os.utime(
                            str(link_file),
                            (Utils.safe_timestamp(created_date).timestamp(),
                             Utils.safe_timestamp(created_date).timestamp()),
                            follow_symlinks=False
                        )

            except FileExistsError:
                log.error('bad link to %s', full_file_name)

        log.warning("Created %d new album folder links", count)
<commit_msg>Update GoogleAlbumsSync.py<commit_after>import logging
import os.path
import shutil
from datetime import datetime
from pathlib import Path
from typing import Dict, Callable

from . import Checks
from . import Utils
from .GoogleAlbumMedia import GoogleAlbumMedia
from .GoogleAlbumsRow import GoogleAlbumsRow
from .GooglePhotosMedia import GooglePhotosMedia
from .GooglePhotosRow import GooglePhotosRow
from .LocalData import LocalData
from .restclient import RestClient

log = logging.getLogger(__name__)

PAGE_SIZE = 100
ALBUM_ITEMS = 50

class GoogleAlbumsSync(object):
    """A Class for managing the indexing and download Google of Albums
    """

    def __init__(self, api: RestClient, root_folder: Path, db: LocalData,
                 flush: bool, photos_path: Path,
                 albums_path: Path, use_flat_path=True, omit_album_date=True,
                 use_hardlinks=False):
        """
        Parameters:
            root_folder: path to the root of local file synchronization
            api: object representing the Google REST API
            db: local database for indexing
            :param use_flat_path:
            :param omit_album_date:
            :param use_hardlinks:
            :param photos_path:
            :param albums_path:
        """
        self._root_folder: Path = root_folder
        self._photos_folder = Path(photos_path)
        self._albums_folder = Path(albums_path)
        self._links_root = self._root_folder / self._albums_folder
        self._photos_root = self._root_folder / self._photos_folder
        self._db: LocalData = db
        self._api: RestClient = api
        self._use_flat_path = use_flat_path
        self._omit_album_date = omit_album_date
        self._use_hardlinks = use_hardlinks
        self.flush = flush
        # these properties are set after construction
        self.album = None
        self.shared_albums = True
        self.album_index = True
        self.use_start_date = False
        self.favourites = False

    @classmethod
    def make_search_parameters(cls, album_id: str,
                               page_token: str = None) -> Dict:
        body = {
            'pageToken': page_token,
            'albumId': album_id,
            'pageSize': PAGE_SIZE
        }
        return body

    def fetch_album_contents(self, album_id: str,
                             add_media_items: bool) -> (datetime, datetime):
        first_date = Utils.maximum_date()
        last_date = Utils.MINIMUM_DATE
        body = self.make_search_parameters(album_id=album_id)
        response = self._api.mediaItems.search.execute(body)
        position = -1
        while response:
            items_json = response.json()
            media_json = items_json.get('mediaItems')
            # cope with empty albums
            if not media_json:
                if not items_json.get('nextPageToken'):
                    break
                else:
                    media_json = []
                    log.warning('*** Empty Media JSON with a Next Page Token')

            for media_item_json in media_json:
                position += 1
                media_item = GooglePhotosMedia(media_item_json)
                log.debug('----%s', media_item.filename)
                self._db.put_album_file(album_id, media_item.id, position)
                last_date = max(media_item.create_date, last_date)
                first_date = min(media_item.create_date, first_date)

                # this adds other users photos from shared albums
                # Todo - This will cause two copies of a file to appear for
                #  those shared items you have imported into your own library.
                #  They will have different RemoteIds, one will point to your
                #  library copy (you own) and one to the shared item in the
                #  the folder. Currently with the meta data available it would
                #  be impossible to eliminate these without eliminating other
                #  cases where date and filename (TITLE) match
                if add_media_items:
                    media_item.set_path_by_date(self._photos_folder,
                                                self._use_flat_path)
                    (num, row) = self._db.file_duplicate_no(
                        str(media_item.filename),
                        str(media_item.relative_folder),
                        media_item.id)
                    # we just learned if there were any duplicates in the db
                    media_item.duplicate_number = num

                    log.debug('Adding album media item %s %s %s',
                          media_item.relative_path, media_item.filename,
                          media_item.duplicate_number)
                    self._db.put_row(
                        GooglePhotosRow.from_media(media_item), False)

            next_page = items_json.get('nextPageToken')
            if next_page:
                body = self.make_search_parameters(album_id=album_id,
                                                   page_token=next_page)
                response = self._api.mediaItems.search.execute(body)
            else:
                break
        return first_date, last_date

    def index_album_media(self):
        # we now index all contents of non-shared albums due to the behaviour
        # reported here https://github.com/gilesknap/gphotos-sync/issues/89
        if self.shared_albums:
            self.index_albums_type(self._api.sharedAlbums.list.execute,
                                   'sharedAlbums', "Shared (titled) Albums",
                                   False, True)
        self.index_albums_type(self._api.albums.list.execute,
                               'albums', "Albums", True, self.album_index)

    def index_albums_type(self, api_function: Callable, item_key: str,
                          description: str, allow_null_title: bool,
                          add_media_items: bool):
        """
        query google photos interface for a list of all albums and index their
        contents into the db
        """
        log.warning('Indexing {} ...'.format(description))

        # when only looking for favourites do not download album contents
        if self.favourites:
            add_media_items = False

        # there are no filters in album listing at present so it always a
        # full rescan - it's quite quick
        count = 0
        response = api_function(pageSize=ALBUM_ITEMS)
        while response:
            results = response.json()
            for album_json in results.get(item_key, []):
                count += 1

                album = GoogleAlbumMedia(album_json)
                indexed_album = self._db.get_album(album_id=album.id)
                already_indexed = indexed_album.size == album.size if \
                    indexed_album else False

                if self.album and self.album != album.orig_name:
                    log.debug('Skipping Album: %s, photos: %d '
                              '(does not match --album)', album.filename,
                              album.size)
                elif not allow_null_title and album.description == 'none':
                    log.debug('Skipping no-title album, photos: %d',
                              album.size)
                elif already_indexed and not self.flush:
                    log.debug('Skipping Album: %s, photos: %d', album.filename,
                              album.size)
                else:
                    log.info('Indexing Album: %s, photos: %d', album.filename,
                             album.size)
                    first_date, last_date = self.fetch_album_contents(
                        album.id, add_media_items)
                    # write the album data down now we know the contents'
                    # date range
                    gar = GoogleAlbumsRow.from_parm(
                        album.id, album.filename, album.size,
                        first_date, last_date)
                    self._db.put_row(gar, update=indexed_album)

            next_page = results.get('nextPageToken')
            if next_page:
                response = api_function(pageSize=ALBUM_ITEMS,
                                        pageToken=next_page)
            else:
                break
        log.warning('Indexed %d %s', count, description)

    def album_folder_name(
        self, album_name: str, start_date: datetime, end_date: datetime
    ) -> Path:
        if self._omit_album_date:
            rel_path = album_name
        else:
            if self.use_start_date:
                d = start_date
            else:
                d = end_date
            year = Utils.safe_str_time(d, '%Y')
            month = Utils.safe_str_time(d, '%m%d')

            if self._use_flat_path:
                rel_path = u"{0}-{1} {2}".format(year, month, album_name)
            else:
                rel_path = Path(year) / u"{0} {1}".format(month, album_name)
        link_folder: Path = self._links_root / rel_path
        return link_folder

    def create_album_content_links(self):
        log.warning("Creating album folder links to media ...")
        count = 0
        album_item = 0
        current_rid = ''

        # always re-create all album links - it is quite fast and a good way
        # to ensure consistency
        # especially now that we have --album-date-by-first-photo
        if self._links_root.exists():
            log.debug('removing previous album links tree')
            shutil.rmtree(self._links_root)
        re_download = not self._links_root.exists()

        for (
                path, file_name, album_name,
                start_date_str, end_date_str,
                rid, created
        ) in \
                self._db.get_album_files(download_again=re_download):
            if current_rid == rid:
                album_item += 1
            else:
                self._db.put_album_downloaded(rid)
                current_rid = rid
                album_item = 0
            end_date = Utils.string_to_date(end_date_str)
            start_date = Utils.string_to_date(start_date_str)

            if len(str(self._root_folder / path)) > Checks.MAX_PATH_LENGTH:
                max_path_len = Checks.MAX_PATH_LENGTH - len(
                    str(self._root_folder)
                )
                log.debug('This path needs to be shrunk: %s' %
                          Path(self._root_folder / path))
                path = path[:max_path_len]
                log.debug('Shrunk to: %s' %
                          Path(self._root_folder / path))

            file_name = file_name[:Checks.MAX_FILENAME_LENGTH]

            full_file_name = self._root_folder / path / file_name

            link_folder: Path = self.album_folder_name(
                album_name, start_date, end_date
            )

            link_file = link_folder / "{:04d}_{}".format(album_item, file_name)
            # incredibly, pathlib.Path.relative_to cannot handle
            # '../' in a relative path !!! reverting to os.path
            relative_filename = os.path.relpath(full_file_name,
                                                str(link_folder))
            log.debug('adding album link %s -> %s', relative_filename,
                      link_file)
            try:
                if not link_folder.is_dir():
                    log.debug('new album folder %s', link_folder)
                    link_folder.mkdir(parents=True)

                created_date = Utils.string_to_date(created)
                if self._use_hardlinks:
                    if full_file_name.exists():
                        os.link(full_file_name, link_file)
                    else:
                        log.debug('skip hardlink for %s, not downloaded',
                                  file_name)
                else:
                    link_file.symlink_to(relative_filename)

                if link_file.exists():
                    count += 1
                    # Windows tries to follow symlinks even though we specify
                    # follow_symlinks=False. So disable setting of link date
                    # if follow not supported
                    if os.utime in os.supports_follow_symlinks:
                        os.utime(
                            str(link_file),
                            (Utils.safe_timestamp(created_date).timestamp(),
                             Utils.safe_timestamp(created_date).timestamp()),
                            follow_symlinks=False
                        )

            except FileExistsError:
                log.error('bad link to %s', full_file_name)

        log.warning("Created %d new album folder links", count)
<|endoftext|>


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

""" Bioinformatics tools developed at the Max Planck Institute for Biology of Ageing"""

import pandas as pd
import numpy as np
import os
import csv
import sys
from suds.client import Client as sudsclient
import os
import ssl

def readGTF(infile):
    """
    Reads a GTF file and labels the respective columns in agreement with GTF file standards:
    'seqname','source','feature','start','end','score','strand','frame','attribute'.

    :param infile: path/to/file.gtf
    :returns: a Pandas dataframe of the respective GTF

    """
    df=pd.read_table(infile, sep='\t', comment="#", header=None, dtype=str)
    df.columns=['seqname','source','feature','start','end','score','strand','frame','attribute']
    #df = df.astype(str) # from DTC
    return df

def retrieve_GTF_field(field,gtf):
    """
    Returns a field of choice from the attribute column of the GTF

    :param field: field to be retrieved
    :returns: a Pandas dataframe with one columns containing the field of choice

    """
    label=field
    field = pd.DataFrame(gtf['attribute'].str.split(field).tolist())[1]
    field = field.astype(str)
    field = pd.DataFrame(field.str.split(';',1).tolist())
    field = pd.DataFrame(field[0].str.split('"').tolist())[1]
    field = pd.DataFrame(field)
    field.columns=[label]
    return field

def attributesGTF(inGTF):
    """
    List the type of attributes in a the attribute section of a GTF file

    :param inGTF: GTF dataframe to be analysed
    :returns: a list of attributes present in the attribute section

    """
    df=pd.DataFrame(inGTF['attribute'].str.split(";").tolist())
    desc=[]
    for i in df.columns.tolist():
        val=df[[i]].dropna()
        val=pd.DataFrame(val[i].str.split(' "').tolist())[0]
        val=list(set(val))
        for v in val:
            if len(v) > 0:
                l=v.split(" ")
                if len(l)>1:
                    l=l[1]
                else:
                    l=l[0]
                desc.append(l)
    desc=list(set(desc))
    finaldesc=[]
    for d in desc:
        if len(d) > 0:
            finaldesc.append(d)
    return finaldesc

def parseGTF(inGTF):
    """
    Reads an extracts all attributes in the attributes section of a GTF and constructs a new dataframe wiht one collumn per attribute instead of the attributes column

    :param inGTF: GTF dataframe to be parsed
    :returns: a dataframe of the orignal input GTF with attributes parsed.

    """

    desc=attributesGTF(inGTF)
    df=inGTF.drop(['attribute'],axis=1)
    for d in desc:
        field=retrieve_GTF_field(d,inGTF)
        df=pd.concat([df,field],axis=1)
    return df

def writeGTF(inGTF,file_path):
    """
    Write a GTF dataframe into a file

    :param inGTF: GTF dataframe to be written. It should either have 9 columns with the last one being the "attributes" section or more than 9 columns where all columns after the 8th will be colapsed into one.
    :param file_path: path/to/the/file.gtf
    :returns: nothing
    """
    cols=inGTF.columns.tolist()
    if len(cols) == 9:
        if 'attribute' in cols:
            df=inGTF
    else:
        df=inGTF[cols[:8]]
        df['attribute']=""
        for c in cols[8:]:
            if c == cols[len(cols)-1]:
                df['attribute']=df['attribute']+c+' "'+inGTF[c].astype(str)+'";'
            else:
                df['attribute']=df['attribute']+c+' "'+inGTF[c].astype(str)+'"; '
    df.to_csv(file_path, sep="\t",header=None,index=None,quoting=csv.QUOTE_NONE)

def GTFtoBED(inGTF,name):
    """
    Transform a GTF dataframe into a bed dataframe

    :param inGTF: GTF dataframe for transformation
    :param name: field of the GTF data frame to be use for the bed 'name' positon

    returns: a bed dataframe with the corresponding bed fiels: 'chrom','chromStart','chromEnd','name','score','strand'
    """
    if name not in inGTF.columns.tolist():
        field=retrieve_GTF_field(name, inGTF)
        inGTF=pd.concat([inGTF,field],axis=1)
    bed=inGTF[['seqname','start','end',name,'score','strand']]
    bed.columns=['chrom','chromStart','chromEnd','name','score','strand']
    bed.drop_duplicates(inplace=True)
    bed.reset_index(inplace=True,drop=True)
    return bed

def writeBED(inBED, file_path):
    """
    Writes a bed dataframe into a bed file.

    :param inBED: bed dataframe to be written.
    :param file_path: /path/to/file.bed

    :returns: nothing

    """
    inBED.to_csv(file_path,index=None,sep="\t",header=None)

david_categories = [
  'GOTERM_BP_FAT', 'GOTERM_CC_FAT', 'GOTERM_MF_FAT', 'KEGG_PATHWAY',
  'BIOCARTA', 'PFAM', 'PROSITE' ]

david_fields = [
  'categoryName', 'termName', 'listHits', 'percent',
  'ease', 'geneIds', 'listTotals', 'popHits', 'popTotals',
  'foldEnrichment', 'bonferroni', 'benjamini', 'afdr']
# include:
# 'fisher'
# 'termName' to 'term' and 'term_name'

def DAVIDenrich(database, categories, user, ids, ids_bg = None, name = '', name_bg = '', verbose = False, p = 0.1, n = 2):

    """
    Queries the DAVID database for an enrichment analysis

    :param database: A string for the database to query, e.g. 'WORMBASE_GENE_ID'
    :param categories: A comma separated string with databases
    :param user: A user ID registered at DAVID for querying
    :param ids: A list with identifiers
    :param name: A string with the name for the query set
    :param ids_bg: A list with the background identifiers to enrich against,
      'None' for whole set
    :param name_bg: A string with the name for the background set
    :param p: Maximum p value for enrichment of a term
    :param n: Minimum number of genes within a term
    :param ct: Maybe another threshold

    :returns: None if no ids match the queried database, or a list with results
    """

    ids = ','.join(ids)
    use_bg = 0
    if ids_bg is not None:
      ids_bg = ','.join(ids_bg)
    ssl._create_default_https_context = ssl._create_unverified_context
    client = sudsclient('https://david.ncifcrf.gov/webservice/services/DAVIDWebService?wsdl')
    client.wsdl.services[0].setlocation('https://david.ncifcrf.gov/webservice/services/DAVIDWebService.DAVIDWebServiceHttpSoap11Endpoint/')
    client_auth = client.service.authenticate(user)
    if verbose:
      print 'User Authentication:', client_auth
      sys.stdout.flush()
    size = client.service.addList(ids, database, name, 0) #| inputListIds,idType,listName,listType)
    if verbose:
      print 'Mapping rate of ids: ', str(size)
      sys.stdout.flush()
    if not float(size) > float(0):
      return None
    if ids_bg is not None:
      size_bg = cliet.service.addList(ids, database, name_bg, 1)
      if verbose:
        print 'Mapping rate of background ids: ', str(size_bg)
        sys.stdout.flush()
    client_categories = client.service.setCategories(categories)
    if verbose:
      print 'Categories used: ', client_categories
      sys.stdout.flush()
    client_report = client.service.getChartReport(p, n)
    size_report = len(client_report)
    if verbose:
      print 'Records reported: ', str(size_report)
      sys.stdout.flush()
    return client_report

def write_DAVID(report, path, sep = '\t'):

  """
  Write the DAVIDenrich list to a file

  :param report: A report list as returned from DAVIDenrich
  :param path: A file path to write to
  :param sep: A character used as separator

  :returns: None
  """

  n = len(report)
  if n < 1:
    print 'warning: report of length 0'
    #sys.exit()
  with open(path, 'w') as out:
    out.write(sep.join(david_fields) + '\n')
    for r in report:
      d = dict(r)
      line = []
      for f in david_fields:
        line.append(str(d[f]))
      out.write(sep.join(line) + '\n')

def pandas_DAVID(report):

  """
  Return a pandas data frame from a DAVIDenrich

  :param report: A report list as returned from DAVIDenrich

  :returns: A pandas data frame
  """

  df = []
  df.append(david_fields)
  for r in report:
    d = dict(r)
    line = []
    for f in david_fields:
      line.append(str(d[f]))
    df.append(line)
  pdf = pd.DataFrame(df)
  return pdf

def getFileFormat (path):

  """
  Return the file format

  :params path: The path to the file

  :returns: None, if file is missing, else one of the strings 'xlsx', 'xls', 'txt'
  """

  xlsx_sign = b'\x50\x4B\x05\06'
  xls_sign = b'\x09\x08\x10\x00\x00\x06\x05\x00'
  ret = None
  signs = [(0, 512, 8), (2, -22, 4)] # whence, offset, size
  for w, o, s in signs:
    with open(path, 'rb') as f:
      f.seek(o, w)
      b = f.read(s)
      if b == xls_sig:
        ret = 'xls'
      elif b == xlsx_sig:
        ret = 'xlsx'
      else
        ret = 'txt'
  return ret

def readDataFrame (path, sheet = None, sep = '\t'):

  """
  Returns a pandas data frame

  :params path: The path to the file
  :params sheet: Sheet name or integer 0-based index for xls[x] files, None for all
  :params sep: A separator for text format

  :returns: A pandas data frame
  """

  ff = file_format(path)
  if ff is None:
    print 'error: file not matching format xls[x]/txt'
    sys.exit()
  if ff in ['xls', 'xlsx']:
    df = pd.read_excel(path, sheetname = sheet)
  else
    df = pd.read_table(path, sep = sep)
  return df

if __name__ == '__main__':
    print "AGEpy"

<commit_msg>bugfix readDataFrame<commit_after>#!/usr/bin/env python

""" Bioinformatics tools developed at the Max Planck Institute for Biology of Ageing"""

import pandas as pd
import numpy as np
import os
import csv
import sys
from suds.client import Client as sudsclient
import os
import ssl

def readGTF(infile):
    """
    Reads a GTF file and labels the respective columns in agreement with GTF file standards:
    'seqname','source','feature','start','end','score','strand','frame','attribute'.

    :param infile: path/to/file.gtf
    :returns: a Pandas dataframe of the respective GTF

    """
    df=pd.read_table(infile, sep='\t', comment="#", header=None, dtype=str)
    df.columns=['seqname','source','feature','start','end','score','strand','frame','attribute']
    #df = df.astype(str) # from DTC
    return df

def retrieve_GTF_field(field,gtf):
    """
    Returns a field of choice from the attribute column of the GTF

    :param field: field to be retrieved
    :returns: a Pandas dataframe with one columns containing the field of choice

    """
    label=field
    field = pd.DataFrame(gtf['attribute'].str.split(field).tolist())[1]
    field = field.astype(str)
    field = pd.DataFrame(field.str.split(';',1).tolist())
    field = pd.DataFrame(field[0].str.split('"').tolist())[1]
    field = pd.DataFrame(field)
    field.columns=[label]
    return field

def attributesGTF(inGTF):
    """
    List the type of attributes in a the attribute section of a GTF file

    :param inGTF: GTF dataframe to be analysed
    :returns: a list of attributes present in the attribute section

    """
    df=pd.DataFrame(inGTF['attribute'].str.split(";").tolist())
    desc=[]
    for i in df.columns.tolist():
        val=df[[i]].dropna()
        val=pd.DataFrame(val[i].str.split(' "').tolist())[0]
        val=list(set(val))
        for v in val:
            if len(v) > 0:
                l=v.split(" ")
                if len(l)>1:
                    l=l[1]
                else:
                    l=l[0]
                desc.append(l)
    desc=list(set(desc))
    finaldesc=[]
    for d in desc:
        if len(d) > 0:
            finaldesc.append(d)
    return finaldesc

def parseGTF(inGTF):
    """
    Reads an extracts all attributes in the attributes section of a GTF and constructs a new dataframe wiht one collumn per attribute instead of the attributes column

    :param inGTF: GTF dataframe to be parsed
    :returns: a dataframe of the orignal input GTF with attributes parsed.

    """

    desc=attributesGTF(inGTF)
    df=inGTF.drop(['attribute'],axis=1)
    for d in desc:
        field=retrieve_GTF_field(d,inGTF)
        df=pd.concat([df,field],axis=1)
    return df

def writeGTF(inGTF,file_path):
    """
    Write a GTF dataframe into a file

    :param inGTF: GTF dataframe to be written. It should either have 9 columns with the last one being the "attributes" section or more than 9 columns where all columns after the 8th will be colapsed into one.
    :param file_path: path/to/the/file.gtf
    :returns: nothing
    """
    cols=inGTF.columns.tolist()
    if len(cols) == 9:
        if 'attribute' in cols:
            df=inGTF
    else:
        df=inGTF[cols[:8]]
        df['attribute']=""
        for c in cols[8:]:
            if c == cols[len(cols)-1]:
                df['attribute']=df['attribute']+c+' "'+inGTF[c].astype(str)+'";'
            else:
                df['attribute']=df['attribute']+c+' "'+inGTF[c].astype(str)+'"; '
    df.to_csv(file_path, sep="\t",header=None,index=None,quoting=csv.QUOTE_NONE)

def GTFtoBED(inGTF,name):
    """
    Transform a GTF dataframe into a bed dataframe

    :param inGTF: GTF dataframe for transformation
    :param name: field of the GTF data frame to be use for the bed 'name' positon

    returns: a bed dataframe with the corresponding bed fiels: 'chrom','chromStart','chromEnd','name','score','strand'
    """
    if name not in inGTF.columns.tolist():
        field=retrieve_GTF_field(name, inGTF)
        inGTF=pd.concat([inGTF,field],axis=1)
    bed=inGTF[['seqname','start','end',name,'score','strand']]
    bed.columns=['chrom','chromStart','chromEnd','name','score','strand']
    bed.drop_duplicates(inplace=True)
    bed.reset_index(inplace=True,drop=True)
    return bed

def writeBED(inBED, file_path):
    """
    Writes a bed dataframe into a bed file.

    :param inBED: bed dataframe to be written.
    :param file_path: /path/to/file.bed

    :returns: nothing

    """
    inBED.to_csv(file_path,index=None,sep="\t",header=None)

david_categories = [
  'GOTERM_BP_FAT', 'GOTERM_CC_FAT', 'GOTERM_MF_FAT', 'KEGG_PATHWAY',
  'BIOCARTA', 'PFAM', 'PROSITE' ]

david_fields = [
  'categoryName', 'termName', 'listHits', 'percent',
  'ease', 'geneIds', 'listTotals', 'popHits', 'popTotals',
  'foldEnrichment', 'bonferroni', 'benjamini', 'afdr']
# include:
# 'fisher'
# 'termName' to 'term' and 'term_name'

def DAVIDenrich(database, categories, user, ids, ids_bg = None, name = '', name_bg = '', verbose = False, p = 0.1, n = 2):

    """
    Queries the DAVID database for an enrichment analysis

    :param database: A string for the database to query, e.g. 'WORMBASE_GENE_ID'
    :param categories: A comma separated string with databases
    :param user: A user ID registered at DAVID for querying
    :param ids: A list with identifiers
    :param name: A string with the name for the query set
    :param ids_bg: A list with the background identifiers to enrich against,
      'None' for whole set
    :param name_bg: A string with the name for the background set
    :param p: Maximum p value for enrichment of a term
    :param n: Minimum number of genes within a term
    :param ct: Maybe another threshold

    :returns: None if no ids match the queried database, or a list with results
    """

    ids = ','.join(ids)
    use_bg = 0
    if ids_bg is not None:
      ids_bg = ','.join(ids_bg)
    ssl._create_default_https_context = ssl._create_unverified_context
    client = sudsclient('https://david.ncifcrf.gov/webservice/services/DAVIDWebService?wsdl')
    client.wsdl.services[0].setlocation('https://david.ncifcrf.gov/webservice/services/DAVIDWebService.DAVIDWebServiceHttpSoap11Endpoint/')
    client_auth = client.service.authenticate(user)
    if verbose:
      print 'User Authentication:', client_auth
      sys.stdout.flush()
    size = client.service.addList(ids, database, name, 0) #| inputListIds,idType,listName,listType)
    if verbose:
      print 'Mapping rate of ids: ', str(size)
      sys.stdout.flush()
    if not float(size) > float(0):
      return None
    if ids_bg is not None:
      size_bg = cliet.service.addList(ids, database, name_bg, 1)
      if verbose:
        print 'Mapping rate of background ids: ', str(size_bg)
        sys.stdout.flush()
    client_categories = client.service.setCategories(categories)
    if verbose:
      print 'Categories used: ', client_categories
      sys.stdout.flush()
    client_report = client.service.getChartReport(p, n)
    size_report = len(client_report)
    if verbose:
      print 'Records reported: ', str(size_report)
      sys.stdout.flush()
    return client_report

def write_DAVID(report, path, sep = '\t'):

  """
  Write the DAVIDenrich list to a file

  :param report: A report list as returned from DAVIDenrich
  :param path: A file path to write to
  :param sep: A character used as separator

  :returns: None
  """

  n = len(report)
  if n < 1:
    print 'warning: report of length 0'
    #sys.exit()
  with open(path, 'w') as out:
    out.write(sep.join(david_fields) + '\n')
    for r in report:
      d = dict(r)
      line = []
      for f in david_fields:
        line.append(str(d[f]))
      out.write(sep.join(line) + '\n')

def pandas_DAVID(report):

  """
  Return a pandas data frame from a DAVIDenrich

  :param report: A report list as returned from DAVIDenrich

  :returns: A pandas data frame
  """

  df = []
  df.append(david_fields)
  for r in report:
    d = dict(r)
    line = []
    for f in david_fields:
      line.append(str(d[f]))
    df.append(line)
  pdf = pd.DataFrame(df)
  return pdf

def getFileFormat (path):

  """
  Return the file format

  :params path: The path to the file

  :returns: None, if file is missing, else one of the strings 'xlsx', 'xls', 'txt'
  """

  xlsx_sign = b'\x50\x4B\x05\06'
  xls_sign = b'\x09\x08\x10\x00\x00\x06\x05\x00'
  ret = None
  signs = [(0, 512, 8), (2, -22, 4)] # whence, offset, size
  for w, o, s in signs:
    with open(path, 'rb') as f:
      f.seek(o, w)
      b = f.read(s)
      if b == xls_sign:
        ret = 'xls'
      elif b == xlsx_sign:
        ret = 'xlsx'
      else:
        ret = 'txt'
  return ret

def readDataFrame (path, sheet = None, sep = '\t'):

  """
  Returns a pandas data frame

  :params path: The path to the file
  :params sheet: Sheet name or integer 0-based index for xls[x] files, None for all
  :params sep: A separator for text format

  :returns: A pandas data frame
  """

  ff = getFileFormat(path)
  if ff is None:
    print 'error: file not matching format xls[x]/txt'
    sys.exit()
  if ff in ['xls', 'xlsx']:
    df = pd.read_excel(path, sheetname = sheet)
  else:
    df = pd.read_table(path, sep = sep)
  return df

if __name__ == '__main__':
    print "AGEpy"

<|endoftext|>