<commit_before>"""
Filename:     plot_wind_stress_metric.py
Author:       Damien Irving, irving.damien@gmail.com
Description:  Plot wind stress metrics

"""

# Import general Python modules

import sys, os, pdb
import argparse
import numpy
import iris
import iris.plot as iplt
import matplotlib.pyplot as plt
import seaborn

# Import my modules

cwd = os.getcwd()
repo_dir = '/'
for directory in cwd.split('/')[1:]:
    repo_dir = os.path.join(repo_dir, directory)
    if directory == 'ocean-analysis':
        break

modules_dir = os.path.join(repo_dir, 'modules')
sys.path.append(modules_dir)
try:
    import general_io as gio
    import timeseries
    import convenient_universal as uconv
except ImportError:
    raise ImportError('Must run this script from anywhere within the ocean-analysis git repo')

# Define functions

experiment_colors = {'historical': 'orange',
                     'historicalGHG': 'red',
                     'historicalAA': 'blue',
                     'rcp26': '#16DB65',
                     'rcp45': '#058C42',
                     'rcp60': '#04471C',
                     'rcp85': '#0D2818'}

basins = {'atlantic': 2,
          'pacific': 3,
          'indian': 5,
          'globe': 100}

def load_data(infile, basin_cube, basin_name):
    """Load, temporally aggregate and spatially slice input data"""

    cube = iris.load_cube(infile, 'surface_downward_x_stress')
    cube = timeseries.convert_to_annual(cube)
    experiment = cube.attributes['experiment_id']

    if not basin_name == 'globe':
        if basin_cube:
            ndim = cube.ndim
            basin_array = uconv.broadcast_array(basin_cube.data, [ndim - 2, ndim - 1], cube.shape)
        else:
            basin_array = uconv.create_basin_array(cube)
        cube.data.mask = numpy.where((cube.data.mask == False) & (basin_array == basins[basin_name]), False, True)

    sh_constraint = iris.Constraint(latitude=lambda cell: -30.0 <= cell < 0.0)
    nh_constraint = iris.Constraint(latitude=lambda cell: 0.0 < cell <= 30.0)
    scrit_constraint = iris.Constraint(latitude=lambda cell: -17.0 <= cell < -13.0)
    ncrit_constraint = iris.Constraint(latitude=lambda cell: 13.0 < cell <= 17.0)

    sh_cube = cube.extract(sh_constraint)
    nh_cube = cube.extract(nh_constraint)
    scrit_cube = cube.extract(scrit_constraint)
    ncrit_cube = cube.extract(ncrit_constraint)

    sh_mean = sh_cube.collapsed(['longitude', 'latitude'], iris.analysis.MEAN)
    nh_mean = nh_cube.collapsed(['longitude', 'latitude'], iris.analysis.MEAN)
    scrit_mean = scrit_cube.collapsed(['longitude', 'latitude'], iris.analysis.MEAN)
    ncrit_mean = ncrit_cube.collapsed(['longitude', 'latitude'], iris.analysis.MEAN)

    return sh_mean, nh_mean, scrit_mean, ncrit_mean, experiment

def plot_data(ax, sh_cube, nh_cube, experiment, previous_experiments, basin, crit=False):
    """Plot the various wind stress metrics."""

    plt.sca(ax)

    iplt.plot(nh_cube, color=experiment_colors[experiment], label=experiment+', NH')
    iplt.plot(sh_cube, color=experiment_colors[experiment], label=experiment+', SH', linestyle='--')

    if not experiment in previous_experiments:
        plt.legend(ncol=2)
    plt.ylabel(str(nh_cube.units))
    plt.xlabel('Year')

    if crit:
        plt.title('Critical latitude (15N/S), %s' %(basin))
    else:
        plt.title('Tropics (0-30N/S), %s' %(basin))

def main(inargs):
    """Run the program."""

    time_constraints = {}
    time_constraints['historical'] = gio.get_time_constraint(inargs.hist_time)
    time_constraints['rcp'] = gio.get_time_constraint(inargs.rcp_time)

    if inargs.basin_file:
        basin_cube = iris.load_cube(inargs.basin_file)
    else:
        basin_cube = None

    width=20
    height=21
    fig = plt.figure(figsize=(width, height))

    ax_dict = {}
    ax_dict[('tropics', 'atlantic')] = fig.add_subplot(3, 2, 1)
    ax_dict[('critical', 'atlantic')] = fig.add_subplot(3, 2, 2)
    ax_dict[('tropics', 'pacific')] = fig.add_subplot(3, 2, 3)
    ax_dict[('critical', 'pacific')] = fig.add_subplot(3, 2, 4)
    ax_dict[('tropics', 'globe')] = fig.add_subplot(3, 2, 5)
    ax_dict[('critical', 'globe')] = fig.add_subplot(3, 2, 6)

    previous_experiments = []
    for filenum, infile in enumerate(inargs.infiles):
        for basin in ['atlantic', 'pacific', 'globe']:
            sh_mean, nh_mean, scrit_mean, ncrit_mean, experiment = load_data(infile, basin_cube, basin)
            plot_data(ax_dict[('tropics', basin)], sh_mean, nh_mean, experiment, previous_experiments, basin, crit=False)
            plot_data(ax_dict[('critical', basin)], scrit_mean, ncrit_mean, experiment, previous_experiments, basin, crit=True)
        previous_experiments.append(experiment)

    title = 'Annual Mean Surface Downward X Stress, %s'  %(sh_mean.attributes['model_id'])
    plt.suptitle(title, size='large')
#    plt.subplots_adjust(top=0.90)

    plt.savefig(inargs.outfile, bbox_inches='tight')
    gio.write_metadata(inargs.outfile, file_info={inargs.infiles[-1]: sh_mean.attributes['history']})

if __name__ == '__main__':

    extra_info ="""

author:
    Damien Irving, irving.damien@gmail.com

"""

    description = 'Plot wind stress metrics'
    parser = argparse.ArgumentParser(description=description,
                                     epilog=extra_info,
                                     argument_default=argparse.SUPPRESS,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("infiles", type=str, nargs='*',
                        help="Input wind stress files")
    parser.add_argument("outfile", type=str, help="Output file")

    parser.add_argument("--basin_file", type=str, default=None,
                        help="Input basin file")

    parser.add_argument("--hist_time", type=str, nargs=2, metavar=('START_DATE', 'END_DATE'), default=None,
                        help="Time period [default = all]")
    parser.add_argument("--rcp_time", type=str, nargs=2, metavar=('START_DATE', 'END_DATE'), default=None,
                        help="Time period [default = all]")

    args = parser.parse_args()
    main(args)
<commit_msg>Added comparison or hemisphere options<commit_after>"""
Filename:     plot_wind_stress_metric.py
Author:       Damien Irving, irving.damien@gmail.com
Description:  Plot wind stress metrics

"""

# Import general Python modules

import sys, os, pdb
import argparse
import numpy
import iris
import iris.plot as iplt
import matplotlib.pyplot as plt
import seaborn

# Import my modules

cwd = os.getcwd()
repo_dir = '/'
for directory in cwd.split('/')[1:]:
    repo_dir = os.path.join(repo_dir, directory)
    if directory == 'ocean-analysis':
        break

modules_dir = os.path.join(repo_dir, 'modules')
sys.path.append(modules_dir)
try:
    import general_io as gio
    import timeseries
    import convenient_universal as uconv
except ImportError:
    raise ImportError('Must run this script from anywhere within the ocean-analysis git repo')

# Define functions

experiment_colors = {'historical': 'orange',
                     'historicalGHG': 'red',
                     'historicalAA': 'blue',
                     'rcp26': '#16DB65',
                     'rcp45': '#058C42',
                     'rcp60': '#04471C',
                     'rcp85': '#0D2818'}

basins = {'atlantic': 2,
          'pacific': 3,
          'indian': 5,
          'globe': 100}

def load_data(infile, basin_cube, basin_name):
    """Load, temporally aggregate and spatially slice input data"""

    cube = iris.load_cube(infile, 'surface_downward_x_stress')
    cube = timeseries.convert_to_annual(cube)
    experiment = cube.attributes['experiment_id']
    if experiment == 'historicalMisc':
        experiment = 'historicalAA'

    if not basin_name == 'globe':
        if basin_cube:
            ndim = cube.ndim
            basin_array = uconv.broadcast_array(basin_cube.data, [ndim - 2, ndim - 1], cube.shape)
        else:
            basin_array = uconv.create_basin_array(cube)
        cube.data.mask = numpy.where((cube.data.mask == False) & (basin_array == basins[basin_name]), False, True)

    sh_constraint = iris.Constraint(latitude=lambda cell: -30.0 <= cell < 0.0)
    nh_constraint = iris.Constraint(latitude=lambda cell: 0.0 < cell <= 30.0)
    scrit_constraint = iris.Constraint(latitude=lambda cell: -17.0 <= cell < -13.0)
    ncrit_constraint = iris.Constraint(latitude=lambda cell: 13.0 < cell <= 17.0)

    sh_cube = cube.extract(sh_constraint)
    nh_cube = cube.extract(nh_constraint)
    scrit_cube = cube.extract(scrit_constraint)
    ncrit_cube = cube.extract(ncrit_constraint)

    sh_mean = sh_cube.collapsed(['longitude', 'latitude'], iris.analysis.MEAN)
    nh_mean = nh_cube.collapsed(['longitude', 'latitude'], iris.analysis.MEAN)
    scrit_mean = scrit_cube.collapsed(['longitude', 'latitude'], iris.analysis.MEAN)
    ncrit_mean = ncrit_cube.collapsed(['longitude', 'latitude'], iris.analysis.MEAN)

    return sh_mean, nh_mean, scrit_mean, ncrit_mean, experiment

def plot_data(ax, sh_cube, nh_cube, experiment, previous_experiments, basin, crit=False, plot_type='comparison'):
    """Plot the various wind stress metrics."""

    plt.sca(ax)

    if plot_type == 'comparison':
        data = nh_cube / sh_cube
        iplt.plot(data, color=experiment_colors[experiment], label=experiment)
        ylabel = 'NH / SH'
    else:
        iplt.plot(nh_cube, color=experiment_colors[experiment], label=experiment+', NH')
        iplt.plot(sh_cube, color=experiment_colors[experiment], label=experiment+', SH', linestyle='--')
        ylabel = str(nh_cube.units)

    if not experiment in previous_experiments:
        plt.legend(ncol=2)
    plt.ylabel(ylabel)
    plt.xlabel('Year')

    if crit:
        plt.title('Critical latitude (15N/S), %s' %(basin))
    else:
        plt.title('Tropics (0-30N/S), %s' %(basin))

def main(inargs):
    """Run the program."""

    time_constraints = {}
    time_constraints['historical'] = gio.get_time_constraint(inargs.hist_time)
    time_constraints['rcp'] = gio.get_time_constraint(inargs.rcp_time)

    if inargs.basin_file:
        basin_cube = iris.load_cube(inargs.basin_file)
    else:
        basin_cube = None

    width=20
    height=21
    fig = plt.figure(figsize=(width, height))

    ax_dict = {}
    ax_dict[('tropics', 'atlantic')] = fig.add_subplot(3, 2, 1)
    ax_dict[('critical', 'atlantic')] = fig.add_subplot(3, 2, 2)
    ax_dict[('tropics', 'pacific')] = fig.add_subplot(3, 2, 3)
    ax_dict[('critical', 'pacific')] = fig.add_subplot(3, 2, 4)
    ax_dict[('tropics', 'globe')] = fig.add_subplot(3, 2, 5)
    ax_dict[('critical', 'globe')] = fig.add_subplot(3, 2, 6)

    previous_experiments = []
    for filenum, infile in enumerate(inargs.infiles):
        for basin in ['atlantic', 'pacific', 'globe']:
            sh_mean, nh_mean, scrit_mean, ncrit_mean, experiment = load_data(infile, basin_cube, basin)
            plot_data(ax_dict[('tropics', basin)], sh_mean, nh_mean, experiment, previous_experiments, basin,
                      crit=False, plot_type=inargs.plot_type)
            plot_data(ax_dict[('critical', basin)], scrit_mean, ncrit_mean, experiment, previous_experiments, basin,
                      crit=True, plot_type=inargs.plot_type)
        previous_experiments.append(experiment)

    title = 'Annual Mean Surface Downward X Stress, %s'  %(sh_mean.attributes['model_id'])
    plt.suptitle(title, size='large')
#    plt.subplots_adjust(top=0.90)

    plt.savefig(inargs.outfile, bbox_inches='tight')
    gio.write_metadata(inargs.outfile, file_info={inargs.infiles[-1]: sh_mean.attributes['history']})

if __name__ == '__main__':

    extra_info ="""

author:
    Damien Irving, irving.damien@gmail.com

"""

    description = 'Plot wind stress metrics'
    parser = argparse.ArgumentParser(description=description,
                                     epilog=extra_info,
                                     argument_default=argparse.SUPPRESS,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("infiles", type=str, nargs='*',
                        help="Input wind stress files")
    parser.add_argument("outfile", type=str, help="Output file")

    parser.add_argument("--basin_file", type=str, default=None,
                        help="Input basin file")
    parser.add_argument("--plot_type", type=str, choices=('hemisphere', 'comparison'), default='comparison',
                        help="Type of data to plot")

    parser.add_argument("--hist_time", type=str, nargs=2, metavar=('START_DATE', 'END_DATE'), default=None,
                        help="Time period [default = all]")
    parser.add_argument("--rcp_time", type=str, nargs=2, metavar=('START_DATE', 'END_DATE'), default=None,
                        help="Time period [default = all]")

    args = parser.parse_args()
    main(args)
<|endoftext|>


<commit_before>#!/usr/bin/python3
import sys, os, pty
sys.path.append(os.path.join(os.path.dirname(__file__),'../bindings/python/'))
import serverboards, pexpect, shlex, re, subprocess, random
import urllib.parse as urlparse
from common import *

@serverboards.rpc_method
def ssh_public_key(**kwargs):
    ensure_ID_RSA()
    with open(ID_RSA_PUB, 'r') as fd:
        return fd.read()

def ssh_urlparse(url):
    if not "://" in url:
        url="ssh://%s"%(url)
    u = urlparse.urlparse(url)
    assert u.scheme == 'ssh'
    port = u.port or '22'
    return (u.hostname, str(port))

def get_fingerprint(url, **kwargs):
    if not url:
        return None
    (hostname, port) = ssh_urlparse(url)
    try:
        output = str(subprocess.check_output(["ssh-keyscan", "-p", port, hostname]), 'utf8')
        output=output.strip().split('\n')
        output.sort()
        serverboards.debug(repr(output))
        return '\n'.join(output)
    except:
        import traceback; traceback.print_exc()
        return None

@serverboards.rpc_method
def remote_fingerprint(url="192.168.1.200", options="", **kwargs):
    serverboards.info(repr(kwargs))
    fingerprint=get_fingerprint(url, options)
    if not fingerprint:
        return {
            "fingerprint":"Cant connect to server at <%s>. Set a valid SSH address."%url,
            "className":"grey disabled",
            "toggle":"Select a valid SSH address"
        }
    fingerprintb = bytes(fingerprint, 'utf8')

    sp=subprocess.Popen(["ssh-keygen", "-lf", "/dev/stdin", "-E", "md5"], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    fingerprint_md5, error = sp.communicate(fingerprintb)

    sp=subprocess.Popen(["ssh-keygen", "-lf", "/dev/stdin", "-E", "sha1"], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    fingerprint_sha1, error = sp.communicate(fingerprintb)

    sp=subprocess.Popen(["ssh-keygen", "-lf", "/dev/stdin", "-E", "sha256"], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    fingerprint_sha256, error = sp.communicate(fingerprintb)

    fingerprint_text = ''.join( [str(fingerprint_md5, 'utf8'), "\n", str(fingerprint_sha1, 'utf8'), "\n", str(fingerprint_sha256, 'utf8')] )

    enabled=False
    if os.path.exists(KNOWN_HOSTS_FILE):
        with open(KNOWN_HOSTS_FILE,"r") as fd:
            enabled = fingerprint in fd.read()

    return {
        "fingerprint_orig" : fingerprint,
        "fingerprint": fingerprint_text,
        "toggle" : "Disable host" if enabled else "Enable host",
        "className" : "red" if enabled else "yellow"
    }

@serverboards.rpc_method
def toggle_remote_fingerprint(url=None, status=None, **args):
    fingerprint=get_fingerprint(url)

    enabled=False

    serverboards.debug(repr(status["fingerprint_orig"]))
    serverboards.debug(repr(fingerprint))
    serverboards.debug(repr(status["fingerprint_orig"]==fingerprint))

    if status["fingerprint_orig"]==fingerprint:
        if os.path.exists(KNOWN_HOSTS_FILE):
            with open(KNOWN_HOSTS_FILE,"r") as fd:
                enabled = fingerprint in fd.read()
        if enabled:
            hostname=fingerprint.split(' ')[0]
            done=False
            with open(KNOWN_HOSTS_FILE, "r") as rd:
                with open(KNOWN_HOSTS_FILE+".bak", "w") as wd:
                    for l in rd.readlines():
                        serverboards.debug(repr((l[:10], hostname, l.startswith(hostname))))

                        if not l.startswith(hostname):
                            wd.write(l)
                        else:
                            done=True
            os.rename(KNOWN_HOSTS_FILE+".bak", KNOWN_HOSTS_FILE)
            serverboards.info("Removed SSH known_host fingerprint for url %s"%(url))
            if done:
                return "Fingerprint removed"
            else:
                raise Exception("Could not remove fingerprint. Fingerprint not found. %s"%(hostname))
        else:
            with open(KNOWN_HOSTS_FILE, "a") as fd:
                fd.write(fingerprint)
                fd.write('\n')
            serverboards.info("Added SSH known_host fingerprint for url %s"%(url))
            return "Fingerprint added"
    raise Exception("Fingerprint has changed")

def test():
    print(get_fingerprint("ssh://testurl","User dmoreno\nProxyCommand ssh 192.168.1.200 -W 10.0.3.92:22 -q"))

if __name__=='__main__':
    if len(sys.argv)>1 and sys.argv[1]=="test":
        test()

    serverboards.loop()
<commit_msg>Another fix to add SSH fingerprints<commit_after>#!/usr/bin/python3
import sys, os, pty
sys.path.append(os.path.join(os.path.dirname(__file__),'../bindings/python/'))
import serverboards, pexpect, shlex, re, subprocess, random
import urllib.parse as urlparse
from common import *

@serverboards.rpc_method
def ssh_public_key(**kwargs):
    ensure_ID_RSA()
    with open(ID_RSA_PUB, 'r') as fd:
        return fd.read()

def ssh_urlparse(url):
    if not "://" in url:
        url="ssh://%s"%(url)
    u = urlparse.urlparse(url)
    assert u.scheme == 'ssh'
    port = u.port or '22'
    return (u.hostname, str(port))

def get_fingerprint(url, *args, **kwargs):
    if not url:
        return None
    (hostname, port) = ssh_urlparse(url)
    try:
        output = str(subprocess.check_output(["ssh-keyscan", "-p", port, hostname]), 'utf8')
        output=output.strip().split('\n')
        output.sort()
        serverboards.debug(repr(output))
        return '\n'.join(output)
    except:
        import traceback; traceback.print_exc()
        return None

@serverboards.rpc_method
def remote_fingerprint(url="192.168.1.200", options="", **kwargs):
    serverboards.info(repr(kwargs))
    fingerprint=get_fingerprint(url, options)
    if not fingerprint:
        return {
            "fingerprint":"Cant connect to server at <%s>. Set a valid SSH address."%url,
            "className":"grey disabled",
            "toggle":"Select a valid SSH address"
        }
    fingerprintb = bytes(fingerprint, 'utf8')

    sp=subprocess.Popen(["ssh-keygen", "-lf", "/dev/stdin", "-E", "md5"], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    fingerprint_md5, error = sp.communicate(fingerprintb)

    sp=subprocess.Popen(["ssh-keygen", "-lf", "/dev/stdin", "-E", "sha1"], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    fingerprint_sha1, error = sp.communicate(fingerprintb)

    sp=subprocess.Popen(["ssh-keygen", "-lf", "/dev/stdin", "-E", "sha256"], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    fingerprint_sha256, error = sp.communicate(fingerprintb)

    fingerprint_text = ''.join( [str(fingerprint_md5, 'utf8'), "\n", str(fingerprint_sha1, 'utf8'), "\n", str(fingerprint_sha256, 'utf8')] )

    enabled=False
    if os.path.exists(KNOWN_HOSTS_FILE):
        with open(KNOWN_HOSTS_FILE,"r") as fd:
            enabled = fingerprint in fd.read()

    return {
        "fingerprint_orig" : fingerprint,
        "fingerprint": fingerprint_text,
        "toggle" : "Disable host" if enabled else "Enable host",
        "className" : "red" if enabled else "yellow"
    }

@serverboards.rpc_method
def toggle_remote_fingerprint(url=None, status=None, **args):
    fingerprint=get_fingerprint(url)

    enabled=False

    serverboards.debug(repr(status["fingerprint_orig"]))
    serverboards.debug(repr(fingerprint))
    serverboards.debug(repr(status["fingerprint_orig"]==fingerprint))

    if status["fingerprint_orig"]==fingerprint:
        if os.path.exists(KNOWN_HOSTS_FILE):
            with open(KNOWN_HOSTS_FILE,"r") as fd:
                enabled = fingerprint in fd.read()
        if enabled:
            hostname=fingerprint.split(' ')[0]
            done=False
            with open(KNOWN_HOSTS_FILE, "r") as rd:
                with open(KNOWN_HOSTS_FILE+".bak", "w") as wd:
                    for l in rd.readlines():
                        serverboards.debug(repr((l[:10], hostname, l.startswith(hostname))))

                        if not l.startswith(hostname):
                            wd.write(l)
                        else:
                            done=True
            os.rename(KNOWN_HOSTS_FILE+".bak", KNOWN_HOSTS_FILE)
            serverboards.info("Removed SSH known_host fingerprint for url %s"%(url))
            if done:
                return "Fingerprint removed"
            else:
                raise Exception("Could not remove fingerprint. Fingerprint not found. %s"%(hostname))
        else:
            with open(KNOWN_HOSTS_FILE, "a") as fd:
                fd.write(fingerprint)
                fd.write('\n')
            serverboards.info("Added SSH known_host fingerprint for url %s"%(url))
            return "Fingerprint added"
    raise Exception("Fingerprint has changed")

def test():
    print(get_fingerprint("ssh://testurl","User dmoreno\nProxyCommand ssh 192.168.1.200 -W 10.0.3.92:22 -q"))

if __name__=='__main__':
    if len(sys.argv)>1 and sys.argv[1]=="test":
        test()

    serverboards.loop()
<|endoftext|>


<commit_before># vim: set expandtab sw=4 ts=4:
'''
Collection of tool functions

Copyright (C) 2014 Dieter Adriaenssens <ruleant@users.sourceforge.net>

This file is part of buildtime-trend
<https://github.com/ruleant/buildtime-trend/>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
'''

import copy
import os
from datetime import datetime
from dateutil.parser import parse
from buildtimetrend.settings import get_project_info

def format_timestamp(timestamp):
    '''
    Format a datetime timestamp (UTC) to ISO format (YYYY-MM-DDTHH:MM:SS)

    Parameters :
    - timestamp : timestamp, seconds since epoch
    '''
    timestamp_datetime = datetime.utcfromtimestamp(timestamp)
    return timestamp_datetime.isoformat()

def split_timestamp(timestamp):
    '''
    Split a timestamp in seconds since epoch in all seperate components :
      year, month, day of month, day of week,
      hour (12 and 24 hour), minute, second

    Parameters :
    - timestamp : timestamp, seconds since epoch
    '''
    return split_datetime(datetime.utcfromtimestamp(timestamp))

def split_isotimestamp(isotimestamp):
    '''
    Split a timestamp in isoformat in all seperate components :
      year, month, day of month, day of week,
      hour (12 and 24 hour), minute, second

    Parameters :
    - isotimestamp : timestamp in isoformat YYYY-MM-DDTHH:MM:SS
    '''
    # use dateutil.parser.parse to pare the timestamp
    return split_datetime(parse(isotimestamp))

def split_datetime(timestamp_datetime):
    '''
    Split a timestamp in datetime format in all seperate components :
      year, month, day of month, day of week,
      hour (12 and 24 hour), minute, second

    Parameters :
    - timestamp_datetime : timestamp in datetime class format
    '''
    timestamp_dict = {}
    timestamp_dict["timestamp"] = timestamp_datetime.isoformat()
    timestamp_dict["year"] = timestamp_datetime.strftime("%Y")
    timestamp_dict["month"] = timestamp_datetime.strftime("%m")
    timestamp_dict["month_short_en"] = timestamp_datetime.strftime("%b")
    timestamp_dict["month_full_en"] = timestamp_datetime.strftime("%B")
    timestamp_dict["day_of_month"] = timestamp_datetime.strftime("%d")
    timestamp_dict["day_of_week"] = timestamp_datetime.strftime("%w")
    timestamp_dict["day_of_week_short_en"] = timestamp_datetime.strftime("%a")
    timestamp_dict["day_of_week_full_en"] = timestamp_datetime.strftime("%A")
    timestamp_dict["hour_12"] = timestamp_datetime.strftime("%I")
    timestamp_dict["hour_ampm"] = timestamp_datetime.strftime("%p")
    timestamp_dict["hour_24"] = timestamp_datetime.strftime("%H")
    timestamp_dict["minute"] = timestamp_datetime.strftime("%M")
    timestamp_dict["second"] = timestamp_datetime.strftime("%S")

    return timestamp_dict

def add_project_info_dict(payload):
    '''
    Adds project info to a dictonary
    Param payload: dictonary payload
    '''
    if payload is None and type(payload) is not dict:
        raise TypeError("param payload should be a dictionary")

    payload_as_dict = copy.deepcopy(payload)

    payload_as_dict["buildtime_trend"] = get_project_info()

    return payload_as_dict

def add_project_info_list(payload):
    '''
    Adds project info to a list of dictionaries
    Param payload: list of dictionaries
    '''
    if payload is None and type(payload) is not list:
        raise TypeError("param payload should be a list")

    payload_as_list = []

    # loop over dicts in payload and add project info to each one
    for event_dict in payload:
        payload_as_list.append(add_project_info_dict(event_dict))

    return payload_as_list

def check_file(filename):
    '''
    Checks if a file exists.

    Parameters :
    - filename : file to be checked
    Returns false if file doesn't exist, true if it exists.
    '''
    # load timestamps file
    if not os.path.isfile(filename):
        print 'File doesn\'t exist : {0}'.format(filename)
        return False

    return True
<commit_msg>fix checking if parameter is dict<commit_after># vim: set expandtab sw=4 ts=4:
'''
Collection of tool functions

Copyright (C) 2014 Dieter Adriaenssens <ruleant@users.sourceforge.net>

This file is part of buildtime-trend
<https://github.com/ruleant/buildtime-trend/>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
'''

import copy
import os
from datetime import datetime
from dateutil.parser import parse
from buildtimetrend.settings import get_project_info

def format_timestamp(timestamp):
    '''
    Format a datetime timestamp (UTC) to ISO format (YYYY-MM-DDTHH:MM:SS)

    Parameters :
    - timestamp : timestamp, seconds since epoch
    '''
    timestamp_datetime = datetime.utcfromtimestamp(timestamp)
    return timestamp_datetime.isoformat()

def split_timestamp(timestamp):
    '''
    Split a timestamp in seconds since epoch in all seperate components :
      year, month, day of month, day of week,
      hour (12 and 24 hour), minute, second

    Parameters :
    - timestamp : timestamp, seconds since epoch
    '''
    return split_datetime(datetime.utcfromtimestamp(timestamp))

def split_isotimestamp(isotimestamp):
    '''
    Split a timestamp in isoformat in all seperate components :
      year, month, day of month, day of week,
      hour (12 and 24 hour), minute, second

    Parameters :
    - isotimestamp : timestamp in isoformat YYYY-MM-DDTHH:MM:SS
    '''
    # use dateutil.parser.parse to pare the timestamp
    return split_datetime(parse(isotimestamp))

def split_datetime(timestamp_datetime):
    '''
    Split a timestamp in datetime format in all seperate components :
      year, month, day of month, day of week,
      hour (12 and 24 hour), minute, second

    Parameters :
    - timestamp_datetime : timestamp in datetime class format
    '''
    timestamp_dict = {}
    timestamp_dict["timestamp"] = timestamp_datetime.isoformat()
    timestamp_dict["year"] = timestamp_datetime.strftime("%Y")
    timestamp_dict["month"] = timestamp_datetime.strftime("%m")
    timestamp_dict["month_short_en"] = timestamp_datetime.strftime("%b")
    timestamp_dict["month_full_en"] = timestamp_datetime.strftime("%B")
    timestamp_dict["day_of_month"] = timestamp_datetime.strftime("%d")
    timestamp_dict["day_of_week"] = timestamp_datetime.strftime("%w")
    timestamp_dict["day_of_week_short_en"] = timestamp_datetime.strftime("%a")
    timestamp_dict["day_of_week_full_en"] = timestamp_datetime.strftime("%A")
    timestamp_dict["hour_12"] = timestamp_datetime.strftime("%I")
    timestamp_dict["hour_ampm"] = timestamp_datetime.strftime("%p")
    timestamp_dict["hour_24"] = timestamp_datetime.strftime("%H")
    timestamp_dict["minute"] = timestamp_datetime.strftime("%M")
    timestamp_dict["second"] = timestamp_datetime.strftime("%S")

    return timestamp_dict

def add_project_info_dict(payload):
    '''
    Adds project info to a dictonary
    Param payload: dictonary payload
    '''
    if payload is None or type(payload) is not dict:
        raise TypeError("param payload should be a dictionary")

    payload_as_dict = copy.deepcopy(payload)

    payload_as_dict["buildtime_trend"] = get_project_info()

    return payload_as_dict

def add_project_info_list(payload):
    '''
    Adds project info to a list of dictionaries
    Param payload: list of dictionaries
    '''
    if payload is None and type(payload) is not list:
        raise TypeError("param payload should be a list")

    payload_as_list = []

    # loop over dicts in payload and add project info to each one
    for event_dict in payload:
        payload_as_list.append(add_project_info_dict(event_dict))

    return payload_as_list

def check_file(filename):
    '''
    Checks if a file exists.

    Parameters :
    - filename : file to be checked
    Returns false if file doesn't exist, true if it exists.
    '''
    # load timestamps file
    if not os.path.isfile(filename):
        print 'File doesn\'t exist : {0}'.format(filename)
        return False

    return True
<|endoftext|>


<commit_before># Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import hashlib
import os
import re
import sys
from typing import List

import jax
from jax.experimental.compilation_cache.file_system_cache import FileSystemCache
import jax._src.lib
from jax._src.lib import xla_client
from absl import logging
from typing import Optional

_cache = None

def initialize_cache(path, max_cache_size_bytes=32 * 2**30):
  """Creates a global cache object. Should only be called once per process.

     max_cache_sixe defaults to 32GiB.
  """
  global _cache
  assert _cache == None, f"The cache path has already been initialized to {_cache._path}"
  _cache = FileSystemCache(path, max_cache_size_bytes)
  logging.warning("Initialized persistent compilation cache at %s", path)

def get_executable(xla_computation, compile_options, backend) -> Optional[xla_client.Executable]:
  """Returns the cached executable if present, or None otherwise."""
  assert _cache is not None, "initialize_cache must be called before you can call get_executable()"
  cache_key = get_cache_key(xla_computation, compile_options, backend)
  xla_executable_serialized = _cache.get(cache_key)
  if not xla_executable_serialized:
    return None
  # TODO(skye): xla_computation.get_hlo_module() is the unoptimized HLO but it should
  #be optimized
  xla_executable_deserialized = backend.deserialize_executable(
      xla_executable_serialized,
      xla_computation.get_hlo_module(),
      compile_options)
  return xla_executable_deserialized

def put_executable(xla_computation, compile_options, executable: xla_client.Executable,
                   backend):
  """Adds 'executable' to the cache, possibly evicting older entries."""
  assert _cache is not None, "initialize_cache must be called before you can call put_executable()"
  cache_key = get_cache_key(xla_computation, compile_options, backend)
  serialized_executable = backend.serialize_executable(executable)
  _cache.put(cache_key, serialized_executable)

def _log_cache_key_hash(hash_obj, last_serialized: str):
  if logging.vlog_is_on(1):
    logging.vlog(1, "get_cache_key hash after serializing %s: %s",
                 last_serialized, hash_obj.digest().hex())

def get_cache_key(xla_computation, compile_options, backend) -> str:
  """Creates a hashed string to use as a key to the compilation cache.

     get_cache_key takes in the xla_computation and compile_options of a program and hashes
     all the components into a uniuqe byte string. This byte string is returned as a regular
     string that is 256 characters long.

     Typical return value example:
      '14ac577cdb2ef6d986078b4054cc9893a9a14a16dbb0d8f37b89167c1f1aacdf'
  """
  hash_obj = hashlib.sha256()
  # The HLO op_name metadata sometimes includes Python function pointers,
  # which cause spurious cache misses. Scrub anything that looks like a
  # function pointer. Example op_name metadata:
  #  op_name="jit(s)/custom_jvp_call_jaxpr
  #   [ jvp_jaxpr_thunk=<function _memoize.<locals>.memoized at 0x7f3fa30f0940>\n
  #   num_consts=0 ]"
  # TODO(skye): in theory this could cause us to scrub meaningful binary proto
  # data. Do something more robust.
  if isinstance(xla_computation, str):
    serialized_hlo = xla_computation.encode()  # MLIR module
  else:
    serialized_hlo = xla_computation.as_serialized_hlo_module_proto()
  scrubbed_hlo = re.sub(b" at 0x[a-f0-9]+>", b" at 0x...>", serialized_hlo)
  hash_obj.update(scrubbed_hlo)
  _log_cache_key_hash(hash_obj, "computation")

  _hash_compile_options(hash_obj, compile_options)
  _log_cache_key_hash(hash_obj, "compile_options")

  hash_obj.update(bytes(jax._src.lib.version))
  _log_cache_key_hash(hash_obj, "jax_lib version")

  _hash_platform(hash_obj, backend)
  _log_cache_key_hash(hash_obj, "the backend")

  _hash_xla_flags(hash_obj)
  _log_cache_key_hash(hash_obj, "XLA flags")

  return hash_obj.digest().hex()

def _hash_compile_options(hash_obj, compile_options_obj):
  assert len(dir(compile_options_obj)) == 31,(f"Unexpected number of CompileOption fields: "
                                              f"{len(dir(compile_options_obj))}. This likely: means that an extra "
                                              f"field was added, and this function needs to be updated.")

  if compile_options_obj.argument_layouts is not None:
    map(lambda shape: hash_obj.update(shape.to_serialized_proto()),
        compile_options_obj.argument_layouts)
  _hash_int(hash_obj, compile_options_obj.parameter_is_tupled_arguments)
  _hash_executable_build_options(hash_obj, compile_options_obj.executable_build_options)
  _hash_bool(hash_obj, compile_options_obj.tuple_arguments)
  _hash_int(hash_obj, compile_options_obj.num_replicas)
  _hash_int(hash_obj, compile_options_obj.num_partitions)
  if compile_options_obj.device_assignment is not None:
    hash_obj.update(compile_options_obj.device_assignment.serialize())

def _hash_executable_build_options(hash_obj, executable_obj):
  expected_options = 31
  assert len(dir(executable_obj)) == expected_options, (
        f"Unexpected number of executable_build_options fields: "
        f"{len(dir(executable_obj))}. This likely means that an extra "
        f"field was added, and this function needs to be updated.")
  if executable_obj.result_layout is not None:
    hash_obj.update(executable_obj.result_layout.to_serialized_proto())
  _hash_int(hash_obj, executable_obj.num_replicas)
  _hash_int(hash_obj, executable_obj.num_partitions)
  _hash_debug_options(hash_obj, executable_obj.debug_options)
  if executable_obj.device_assignment is not None:
    hash_obj.update(executable_obj.device_assignment.serialize())
  _hash_bool(hash_obj, executable_obj.use_spmd_partitioning)
  _hash_bool(hash_obj, executable_obj.allow_spmd_sharding_propagation_to_output)

def _hash_debug_options(hash_obj, debug_obj):
  _hash_bool(hash_obj, debug_obj.xla_cpu_enable_fast_math)
  _hash_bool(hash_obj, debug_obj.xla_cpu_fast_math_honor_infs)
  _hash_bool(hash_obj, debug_obj.xla_cpu_fast_math_honor_nans)
  _hash_bool(hash_obj, debug_obj.xla_cpu_fast_math_honor_division)
  _hash_bool(hash_obj, debug_obj.xla_cpu_fast_math_honor_functions)
  _hash_bool(hash_obj, debug_obj.xla_gpu_enable_fast_min_max)
  _hash_int(hash_obj, debug_obj.xla_backend_optimization_level)
  _hash_bool(hash_obj, debug_obj.xla_cpu_enable_xprof_traceme)
  _hash_bool(hash_obj, debug_obj.xla_llvm_disable_expensive_passes)
  _hash_bool(hash_obj, debug_obj.xla_test_all_input_layouts)

def _hash_platform(hash_obj, backend):
  _hash_string(hash_obj, backend.platform)
  _hash_string(hash_obj, backend.platform_version)
  _hash_string(hash_obj, backend.runtime_type)

_xla_flags_to_exclude_from_cache_key = [
    "--xla_dump_compress_protos",
    "--xla_dump_module_metadata",
    "--xla_dump_max_hlo_modules",
    "--xla_dump_include_timestamp",
    "--xla_dump_hlo_pass_re",
    "--xla_dump_hlo_module_re",
    "--xla_dump_hlo_snapshots",
    "--xla_dump_fusion_visualization",
    "--xla_dump_hlo_as_url",
    "--xla_dump_hlo_as_proto",
    "--xla_dump_hlo_as_text",
    "--xla_dump_to",
    "--xla_force_host_platform_device_count",
    "--xla_dump_disable_metadata",
    "--xla_dump_hlo_pipeline_re",
]

extra_flag_prefixes_to_include_in_cache_key: List[str] = []

def _hash_xla_flags(hash_obj):
  xla_flags = []

  xla_flags_env_var = os.getenv("XLA_FLAGS")
  if xla_flags_env_var:
    xla_flags.extend(xla_flags_env_var.split())

  for arg in sys.argv:
    if arg.startswith("--xla") or any(
        arg.startswith(p) for p in extra_flag_prefixes_to_include_in_cache_key):
      xla_flags.append(arg)

  # N.B. all XLA flags that take an argument must use '=' and not a space
  # (e.g. --xla_force_host_platform_device_count=8) (I think).
  for flag in xla_flags:
    if flag.split('=')[0] in _xla_flags_to_exclude_from_cache_key:
      logging.vlog(1, "Not including XLA flag in cache key: %s", flag)
      continue
    logging.vlog(1, "Including XLA flag in cache key: %s", flag)
    _hash_string(hash_obj, flag)

def _hash_int(hash_obj, int_var):
  hash_obj.update(int_var.to_bytes(8, byteorder='big'))

def _hash_bool(hash_obj, bool_var):
  hash_obj.update(bool_var.to_bytes(1, byteorder='big'))

def _hash_string(hash_obj, str_var):
  hash_obj.update(str_var.encode('utf-8').strip())

def is_initialized():
  return _cache is not None
<commit_msg>[jax] Add log message when cache entries are written.<commit_after># Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import hashlib
import os
import re
import sys
from typing import List, Optional

import jax
from jax.experimental.compilation_cache.file_system_cache import FileSystemCache
import jax._src.lib
from jax._src.lib import xla_client
from absl import logging

_cache = None

def initialize_cache(path, max_cache_size_bytes=32 * 2**30):
  """Creates a global cache object. Should only be called once per process.

     max_cache_sixe defaults to 32GiB.
  """
  global _cache
  assert _cache == None, f"The cache path has already been initialized to {_cache._path}"
  _cache = FileSystemCache(path, max_cache_size_bytes)
  logging.warning("Initialized persistent compilation cache at %s", path)

def get_executable(xla_computation, compile_options, backend) -> Optional[xla_client.Executable]:
  """Returns the cached executable if present, or None otherwise."""
  assert _cache is not None, "initialize_cache must be called before you can call get_executable()"
  cache_key = get_cache_key(xla_computation, compile_options, backend)
  xla_executable_serialized = _cache.get(cache_key)
  if not xla_executable_serialized:
    return None
  # TODO(skye): xla_computation.get_hlo_module() is the unoptimized HLO but it should
  #be optimized
  xla_executable_deserialized = backend.deserialize_executable(
      xla_executable_serialized,
      xla_computation.get_hlo_module(),
      compile_options)
  return xla_executable_deserialized

def put_executable(xla_computation, compile_options, executable: xla_client.Executable,
                   backend):
  """Adds 'executable' to the cache, possibly evicting older entries."""
  assert _cache is not None, "initialize_cache must be called before you can call put_executable()"
  cache_key = get_cache_key(xla_computation, compile_options, backend)
  logging.info('Writing %s to persistent compilation cache with key %s.',
               xla_computation.name(), cache_key)
  serialized_executable = backend.serialize_executable(executable)
  _cache.put(cache_key, serialized_executable)

def _log_cache_key_hash(hash_obj, last_serialized: str):
  if logging.vlog_is_on(1):
    logging.vlog(1, "get_cache_key hash after serializing %s: %s",
                 last_serialized, hash_obj.digest().hex())

def get_cache_key(xla_computation, compile_options, backend) -> str:
  """Creates a hashed string to use as a key to the compilation cache.

     get_cache_key takes in the xla_computation and compile_options of a program and hashes
     all the components into a uniuqe byte string. This byte string is returned as a regular
     string that is 256 characters long.

     Typical return value example:
      '14ac577cdb2ef6d986078b4054cc9893a9a14a16dbb0d8f37b89167c1f1aacdf'
  """
  hash_obj = hashlib.sha256()
  # The HLO op_name metadata sometimes includes Python function pointers,
  # which cause spurious cache misses. Scrub anything that looks like a
  # function pointer. Example op_name metadata:
  #  op_name="jit(s)/custom_jvp_call_jaxpr
  #   [ jvp_jaxpr_thunk=<function _memoize.<locals>.memoized at 0x7f3fa30f0940>\n
  #   num_consts=0 ]"
  # TODO(skye): in theory this could cause us to scrub meaningful binary proto
  # data. Do something more robust.
  if isinstance(xla_computation, str):
    serialized_hlo = xla_computation.encode()  # MLIR module
  else:
    serialized_hlo = xla_computation.as_serialized_hlo_module_proto()
  scrubbed_hlo = re.sub(b" at 0x[a-f0-9]+>", b" at 0x...>", serialized_hlo)
  hash_obj.update(scrubbed_hlo)
  _log_cache_key_hash(hash_obj, "computation")

  _hash_compile_options(hash_obj, compile_options)
  _log_cache_key_hash(hash_obj, "compile_options")

  hash_obj.update(bytes(jax._src.lib.version))
  _log_cache_key_hash(hash_obj, "jax_lib version")

  _hash_platform(hash_obj, backend)
  _log_cache_key_hash(hash_obj, "the backend")

  _hash_xla_flags(hash_obj)
  _log_cache_key_hash(hash_obj, "XLA flags")

  return hash_obj.digest().hex()

def _hash_compile_options(hash_obj, compile_options_obj):
  assert len(dir(compile_options_obj)) == 31,(f"Unexpected number of CompileOption fields: "
                                              f"{len(dir(compile_options_obj))}. This likely: means that an extra "
                                              f"field was added, and this function needs to be updated.")

  if compile_options_obj.argument_layouts is not None:
    map(lambda shape: hash_obj.update(shape.to_serialized_proto()),
        compile_options_obj.argument_layouts)
  _hash_int(hash_obj, compile_options_obj.parameter_is_tupled_arguments)
  _hash_executable_build_options(hash_obj, compile_options_obj.executable_build_options)
  _hash_bool(hash_obj, compile_options_obj.tuple_arguments)
  _hash_int(hash_obj, compile_options_obj.num_replicas)
  _hash_int(hash_obj, compile_options_obj.num_partitions)
  if compile_options_obj.device_assignment is not None:
    hash_obj.update(compile_options_obj.device_assignment.serialize())

def _hash_executable_build_options(hash_obj, executable_obj):
  expected_options = 31
  assert len(dir(executable_obj)) == expected_options, (
        f"Unexpected number of executable_build_options fields: "
        f"{len(dir(executable_obj))}. This likely means that an extra "
        f"field was added, and this function needs to be updated.")
  if executable_obj.result_layout is not None:
    hash_obj.update(executable_obj.result_layout.to_serialized_proto())
  _hash_int(hash_obj, executable_obj.num_replicas)
  _hash_int(hash_obj, executable_obj.num_partitions)
  _hash_debug_options(hash_obj, executable_obj.debug_options)
  if executable_obj.device_assignment is not None:
    hash_obj.update(executable_obj.device_assignment.serialize())
  _hash_bool(hash_obj, executable_obj.use_spmd_partitioning)
  _hash_bool(hash_obj, executable_obj.allow_spmd_sharding_propagation_to_output)

def _hash_debug_options(hash_obj, debug_obj):
  _hash_bool(hash_obj, debug_obj.xla_cpu_enable_fast_math)
  _hash_bool(hash_obj, debug_obj.xla_cpu_fast_math_honor_infs)
  _hash_bool(hash_obj, debug_obj.xla_cpu_fast_math_honor_nans)
  _hash_bool(hash_obj, debug_obj.xla_cpu_fast_math_honor_division)
  _hash_bool(hash_obj, debug_obj.xla_cpu_fast_math_honor_functions)
  _hash_bool(hash_obj, debug_obj.xla_gpu_enable_fast_min_max)
  _hash_int(hash_obj, debug_obj.xla_backend_optimization_level)
  _hash_bool(hash_obj, debug_obj.xla_cpu_enable_xprof_traceme)
  _hash_bool(hash_obj, debug_obj.xla_llvm_disable_expensive_passes)
  _hash_bool(hash_obj, debug_obj.xla_test_all_input_layouts)

def _hash_platform(hash_obj, backend):
  _hash_string(hash_obj, backend.platform)
  _hash_string(hash_obj, backend.platform_version)
  _hash_string(hash_obj, backend.runtime_type)

_xla_flags_to_exclude_from_cache_key = [
    "--xla_dump_compress_protos",
    "--xla_dump_module_metadata",
    "--xla_dump_max_hlo_modules",
    "--xla_dump_include_timestamp",
    "--xla_dump_hlo_pass_re",
    "--xla_dump_hlo_module_re",
    "--xla_dump_hlo_snapshots",
    "--xla_dump_fusion_visualization",
    "--xla_dump_hlo_as_url",
    "--xla_dump_hlo_as_proto",
    "--xla_dump_hlo_as_text",
    "--xla_dump_to",
    "--xla_force_host_platform_device_count",
    "--xla_dump_disable_metadata",
    "--xla_dump_hlo_pipeline_re",
]

extra_flag_prefixes_to_include_in_cache_key: List[str] = []

def _hash_xla_flags(hash_obj):
  xla_flags = []

  xla_flags_env_var = os.getenv("XLA_FLAGS")
  if xla_flags_env_var:
    xla_flags.extend(xla_flags_env_var.split())

  for arg in sys.argv:
    if arg.startswith("--xla") or any(
        arg.startswith(p) for p in extra_flag_prefixes_to_include_in_cache_key):
      xla_flags.append(arg)

  # N.B. all XLA flags that take an argument must use '=' and not a space
  # (e.g. --xla_force_host_platform_device_count=8) (I think).
  for flag in xla_flags:
    if flag.split('=')[0] in _xla_flags_to_exclude_from_cache_key:
      logging.vlog(1, "Not including XLA flag in cache key: %s", flag)
      continue
    logging.vlog(1, "Including XLA flag in cache key: %s", flag)
    _hash_string(hash_obj, flag)

def _hash_int(hash_obj, int_var):
  hash_obj.update(int_var.to_bytes(8, byteorder='big'))

def _hash_bool(hash_obj, bool_var):
  hash_obj.update(bool_var.to_bytes(1, byteorder='big'))

def _hash_string(hash_obj, str_var):
  hash_obj.update(str_var.encode('utf-8').strip())

def is_initialized():
  return _cache is not None
<|endoftext|>


<commit_before>import logging
from typing import Union

import flask
from flask import Blueprint
from flask import Response
from flask import flash
from flask import redirect
from flask import render_template
from flask import request
from flask import url_for
from flask_login import login_required, current_user
from pyswagger.contrib.client import flask
from sqlalchemy import asc
from sqlalchemy import or_
from sqlalchemy.orm import joinedload

from waitlist import db
from waitlist.blueprints.settings import add_menu_entry
from waitlist.permissions import perm_manager
from waitlist.permissions.manager import StaticPermissions, StaticRoles
from waitlist.signal.signals import send_account_created, send_roles_changed, send_account_status_change
from waitlist.storage.database import Account, Character, Role, linked_chars, APICacheCharacterInfo
from waitlist.utility.settings import sget_resident_mail, sget_tbadge_mail, sget_other_mail, sget_other_topic, \
    sget_tbadge_topic, sget_resident_topic
from waitlist.utility.utils import get_random_token

from waitlist.utility import outgate

bp = Blueprint('accounts', __name__)
logger = logging.getLogger(__name__)

perm_manager.define_permission('include_in_accountlist')
perm_manager.define_permission('accounts_edit')
perm_manager.define_permission('settings_access')
perm_manager.define_permission('admin')
perm_manager.define_permission('accounts_download_list')

@bp.route("/", methods=["GET", "POST"])
@login_required
@perm_manager.require('accounts_edit')
def accounts():
    if request.method == "POST":
        acc_name = request.form['account_name']

        acc_roles = request.form.getlist('account_roles')

        char_name = request.form['default_char_name']
        char_name = char_name.strip()

        note = request.form['change_note'].strip()

        char_info = outgate.character.get_info_by_name(char_name)
        if char_info is None:
            flash(f"A Character named {char_name} does not exist!")
        else:
            char_id = char_info.id
            acc = Account()
            acc.username = acc_name

            acc.login_token = get_random_token(16)

            if len(acc_roles) > 0:
                db_roles = db.session.query(Role).filter(or_(Role.name == name for name in acc_roles)).all()
                for role in db_roles:
                    acc.roles.append(role)

            db.session.add(acc)

            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            acc.characters.append(character)

            db.session.flush()

            acc.current_char = char_id

            db.session.commit()
            send_account_created(accounts, acc.id, current_user.id, acc_roles, 'Creating account. ' + note)

    roles = db.session.query(Role).order_by(Role.name).all()
    accs = db.session.query(Account).order_by(asc(Account.disabled)).order_by(Account.username).all()
    mails = {
        'resident': [sget_resident_mail(), sget_resident_topic()],
        'tbadge': [sget_tbadge_mail(), sget_tbadge_topic()],
        'other': [sget_other_mail(), sget_other_topic()]
    }

    return render_template("settings/accounts.html", roles=roles, accounts=accs, mails=mails)

@bp.route("/edit", methods=["POST"])
@login_required
@perm_manager.require('accounts_edit')
def account_edit():
    acc_id = int(request.form['account_id'])
    acc_name = request.form['account_name']

    note = request.form['change_note'].strip()

    acc_roles = request.form.getlist('account_roles')

    char_name = request.form['default_char_name']
    char_name = char_name.strip()
    if char_name == "":
        char_name = None

    acc = db.session.query(Account).filter(Account.id == acc_id).first()
    if acc is None:
        return flask.abort(400)

    if acc.username != acc_name:
        acc.username = acc_name

    # if there are roles, add new ones, remove the ones that aren't there
    if len(acc_roles) > 0:
        roles_new = {}
        for r in acc_roles:
            roles_new[r] = True

        # db_roles = session.query(Role).filter(or_(Role.name == name for name in acc_roles)).all()
        roles_to_remove = []
        for role in acc.roles:
            if role.name in roles_new:
                del roles_new[role.name]  # remove because it is already in the db
            else:
                # remove the roles because it not submitted anymore
                # only remove admin if current user is an admin
                if role.name == StaticRoles.ADMIN and not perm_manager.get_permission(StaticPermissions.ADMIN).can():
                    continue
                roles_to_remove.append(role)  # mark for removal

        for role in roles_to_remove:
            acc.roles.remove(role)

        # if it is not an admin remove admin role from new roles
        if not perm_manager.get_permission('admin').can():
            if 'admin' in roles_new:
                del roles_new['admin']

        # add remaining roles
        if len(roles_new) > 0:
            new_db_roles = db.session.query(Role).filter(or_(Role.name == name for name in roles_new))
            for role in new_db_roles:
                acc.roles.append(role)
        if len(roles_new) > 0 or len(roles_to_remove) > 0:
            send_roles_changed(account_edit, acc.id, current_user.id, [x for x in roles_new],
                               [x.name for x in roles_to_remove], note)
    else:
        # make sure all roles are removed
        roles_to_remove = []
        for role in acc.roles:
            # only remove admin if current user is an admin
            if role.name == StaticRoles.ADMIN and not perm_manager.get_permission(StaticPermissions.ADMIN).can():
                continue
            roles_to_remove.append(role)

        if len(roles_to_remove) > 0:
            for role in roles_to_remove:
                acc.roles.remove(role)
            db.session.flush()
            send_roles_changed(account_edit, acc.id, current_user.id, [], [x.name for x in roles_to_remove], note)

    if char_name is not None:
        char_info = outgate.character.get_info_by_name(char_name)
        if char_info is None:
            flash(f"Character with name {char_name} could not be found!")
        else:
            char_id = char_info.id
            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                # lets make sure we have the correct name (case)
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            # check if character is linked to this account
            link = db.session.query(linked_chars) \
                .filter((linked_chars.c.id == acc_id) & (linked_chars.c.char_id == char_id)).first()
            if link is None:
                acc.characters.append(character)

            db.session.flush()
            acc.current_char = char_id

    db.session.commit()
    return redirect(url_for('.accounts'), code=303)

@bp.route("/self_edit", methods=["POST"])
@login_required
@perm_manager.require('settings_access')
def account_self_edit():
    acc_id = current_user.id

    char_name = request.form['default_char_name']
    char_name = char_name.strip()
    if char_name == "":
        char_name = None

    acc = db.session.query(Account).filter(Account.id == acc_id).first()
    if acc is None:
        return flask.abort(400)

    if char_name is not None:
        char_info = outgate.character.get_info_by_name(char_name)

        if char_info is None:
            flash("Character with name {char_name} could not be found!")
        else:
            char_id = char_info.id
            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                # lets make sure we have the correct name (case)
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            # check if character is linked to this account
            link = db.session.query(linked_chars) \
                .filter((linked_chars.c.id == acc_id) & (linked_chars.c.char_id == char_id)).first()
            if link is None:
                acc.characters.append(character)

            db.session.flush()
            if acc.current_char != char_id:
                # remove all the access tokens
                if acc.ssoToken is not None:
                    db.session.delete(acc.ssoToken)
                acc.current_char = char_id

    db.session.commit()
    return redirect(url_for('.account_self'), code=303)

@bp.route("/self", methods=["GET"])
@login_required
@perm_manager.require('settings_access')
def account_self():
    acc = db.session.query(Account).filter(Account.id == current_user.id).first()
    return render_template("settings/self.html", account=acc)

@bp.route("/api/account/disabled", methods=['POST'])
@login_required
@perm_manager.require('settings_access')
def account_disabled():
    accid: int = int(request.form['id'])
    acc: Account = db.session.query(Account).filter(Account.id == accid).first()
    status: Union(str, bool) = request.form['disabled']
    send_account_status_change(account_disabled, acc.id, current_user.id, status)
    logger.info("%s sets account %s to %s", current_user.username, acc.username, status)
    if status == 'false':
        status = False
    else:
        status = True

    acc.disabled = status
    db.session.commit()
    return "OK"

@bp.route("/api/account/<int:acc_id>", methods=["DELETE"])
@login_required
@perm_manager.require('admin')
def api_account_delete(acc_id: int) -> Response:
    db.session.query(Account).filter(Account.id == acc_id).delete()
    db.session.commit()
    return flask.jsonify(status="OK")

@bp.route('/accounts/downloadlist/cvs')
@login_required
@perm_manager.require('accounts_download_list')
def accounts_download_csv() -> Response:
    def iter_accs(data):
        for account in data:
            for ci, char in enumerate(account.characters):
                if ci > 0:
                    yield ", " + char.eve_name
                else:
                    yield char.eve_name
            yield '\n'

    permission = perm_manager.get_permission('include_in_accountlist')
    role_check = not Account.disabled
    for role_need in permission.needs:
        if role_need.method != 'role':
            continue
        role_check &= Role.name == role_need.value

    # noinspection PyPep8
    accs = db.session.query(Account).options(joinedload('characters')).join(Account.roles).filter(
        role_check ).order_by(
        Account.username).all()

    response = Response(iter_accs(accs), mimetype='text/csv')
    response.headers['Content-Disposition'] = 'attachment; filename=accounts.csv'
    return response

add_menu_entry('accounts.accounts', 'Accounts', perm_manager.get_permission('accounts_edit').can)
add_menu_entry('accounts.account_self', 'Own Settings', lambda: True)
<commit_msg>fixed some rules for generating full fc list<commit_after>import logging
from typing import Union

import flask
from flask import Blueprint
from flask import Response
from flask import flash
from flask import redirect
from flask import render_template
from flask import request
from flask import url_for
from flask_login import login_required, current_user
from pyswagger.contrib.client import flask
from sqlalchemy import asc
from sqlalchemy import or_
from sqlalchemy.orm import joinedload

from waitlist import db
from waitlist.blueprints.settings import add_menu_entry
from waitlist.permissions import perm_manager
from waitlist.permissions.manager import StaticPermissions, StaticRoles
from waitlist.signal.signals import send_account_created, send_roles_changed, send_account_status_change
from waitlist.storage.database import Account, Character, Role, linked_chars, APICacheCharacterInfo
from waitlist.utility.settings import sget_resident_mail, sget_tbadge_mail, sget_other_mail, sget_other_topic, \
    sget_tbadge_topic, sget_resident_topic
from waitlist.utility.utils import get_random_token

from waitlist.utility import outgate

bp = Blueprint('accounts', __name__)
logger = logging.getLogger(__name__)

perm_manager.define_permission('include_in_accountlist')
perm_manager.define_permission('accounts_edit')
perm_manager.define_permission('settings_access')
perm_manager.define_permission('admin')
perm_manager.define_permission('accounts_download_list')

@bp.route("/", methods=["GET", "POST"])
@login_required
@perm_manager.require('accounts_edit')
def accounts():
    if request.method == "POST":
        acc_name = request.form['account_name']

        acc_roles = request.form.getlist('account_roles')

        char_name = request.form['default_char_name']
        char_name = char_name.strip()

        note = request.form['change_note'].strip()

        char_info = outgate.character.get_info_by_name(char_name)
        if char_info is None:
            flash(f"A Character named {char_name} does not exist!")
        else:
            char_id = char_info.id
            acc = Account()
            acc.username = acc_name

            acc.login_token = get_random_token(16)

            if len(acc_roles) > 0:
                db_roles = db.session.query(Role).filter(or_(Role.name == name for name in acc_roles)).all()
                for role in db_roles:
                    acc.roles.append(role)

            db.session.add(acc)

            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            acc.characters.append(character)

            db.session.flush()

            acc.current_char = char_id

            db.session.commit()
            send_account_created(accounts, acc.id, current_user.id, acc_roles, 'Creating account. ' + note)

    roles = db.session.query(Role).order_by(Role.name).all()
    accs = db.session.query(Account).order_by(asc(Account.disabled)).order_by(Account.username).all()
    mails = {
        'resident': [sget_resident_mail(), sget_resident_topic()],
        'tbadge': [sget_tbadge_mail(), sget_tbadge_topic()],
        'other': [sget_other_mail(), sget_other_topic()]
    }

    return render_template("settings/accounts.html", roles=roles, accounts=accs, mails=mails)

@bp.route("/edit", methods=["POST"])
@login_required
@perm_manager.require('accounts_edit')
def account_edit():
    acc_id = int(request.form['account_id'])
    acc_name = request.form['account_name']

    note = request.form['change_note'].strip()

    acc_roles = request.form.getlist('account_roles')

    char_name = request.form['default_char_name']
    char_name = char_name.strip()
    if char_name == "":
        char_name = None

    acc = db.session.query(Account).filter(Account.id == acc_id).first()
    if acc is None:
        return flask.abort(400)

    if acc.username != acc_name:
        acc.username = acc_name

    # if there are roles, add new ones, remove the ones that aren't there
    if len(acc_roles) > 0:
        roles_new = {}
        for r in acc_roles:
            roles_new[r] = True

        # db_roles = session.query(Role).filter(or_(Role.name == name for name in acc_roles)).all()
        roles_to_remove = []
        for role in acc.roles:
            if role.name in roles_new:
                del roles_new[role.name]  # remove because it is already in the db
            else:
                # remove the roles because it not submitted anymore
                # only remove admin if current user is an admin
                if role.name == StaticRoles.ADMIN and not perm_manager.get_permission(StaticPermissions.ADMIN).can():
                    continue
                roles_to_remove.append(role)  # mark for removal

        for role in roles_to_remove:
            acc.roles.remove(role)

        # if it is not an admin remove admin role from new roles
        if not perm_manager.get_permission('admin').can():
            if 'admin' in roles_new:
                del roles_new['admin']

        # add remaining roles
        if len(roles_new) > 0:
            new_db_roles = db.session.query(Role).filter(or_(Role.name == name for name in roles_new))
            for role in new_db_roles:
                acc.roles.append(role)
        if len(roles_new) > 0 or len(roles_to_remove) > 0:
            send_roles_changed(account_edit, acc.id, current_user.id, [x for x in roles_new],
                               [x.name for x in roles_to_remove], note)
    else:
        # make sure all roles are removed
        roles_to_remove = []
        for role in acc.roles:
            # only remove admin if current user is an admin
            if role.name == StaticRoles.ADMIN and not perm_manager.get_permission(StaticPermissions.ADMIN).can():
                continue
            roles_to_remove.append(role)

        if len(roles_to_remove) > 0:
            for role in roles_to_remove:
                acc.roles.remove(role)
            db.session.flush()
            send_roles_changed(account_edit, acc.id, current_user.id, [], [x.name for x in roles_to_remove], note)

    if char_name is not None:
        char_info = outgate.character.get_info_by_name(char_name)
        if char_info is None:
            flash(f"Character with name {char_name} could not be found!")
        else:
            char_id = char_info.id
            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                # lets make sure we have the correct name (case)
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            # check if character is linked to this account
            link = db.session.query(linked_chars) \
                .filter((linked_chars.c.id == acc_id) & (linked_chars.c.char_id == char_id)).first()
            if link is None:
                acc.characters.append(character)

            db.session.flush()
            acc.current_char = char_id

    db.session.commit()
    return redirect(url_for('.accounts'), code=303)

@bp.route("/self_edit", methods=["POST"])
@login_required
@perm_manager.require('settings_access')
def account_self_edit():
    acc_id = current_user.id

    char_name = request.form['default_char_name']
    char_name = char_name.strip()
    if char_name == "":
        char_name = None

    acc = db.session.query(Account).filter(Account.id == acc_id).first()
    if acc is None:
        return flask.abort(400)

    if char_name is not None:
        char_info = outgate.character.get_info_by_name(char_name)

        if char_info is None:
            flash("Character with name {char_name} could not be found!")
        else:
            char_id = char_info.id
            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                # lets make sure we have the correct name (case)
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            # check if character is linked to this account
            link = db.session.query(linked_chars) \
                .filter((linked_chars.c.id == acc_id) & (linked_chars.c.char_id == char_id)).first()
            if link is None:
                acc.characters.append(character)

            db.session.flush()
            if acc.current_char != char_id:
                # remove all the access tokens
                if acc.ssoToken is not None:
                    db.session.delete(acc.ssoToken)
                acc.current_char = char_id

    db.session.commit()
    return redirect(url_for('.account_self'), code=303)

@bp.route("/self", methods=["GET"])
@login_required
@perm_manager.require('settings_access')
def account_self():
    acc = db.session.query(Account).filter(Account.id == current_user.id).first()
    return render_template("settings/self.html", account=acc)

@bp.route("/api/account/disabled", methods=['POST'])
@login_required
@perm_manager.require('settings_access')
def account_disabled():
    accid: int = int(request.form['id'])
    acc: Account = db.session.query(Account).filter(Account.id == accid).first()
    status: Union(str, bool) = request.form['disabled']
    send_account_status_change(account_disabled, acc.id, current_user.id, status)
    logger.info("%s sets account %s to %s", current_user.username, acc.username, status)
    if status == 'false':
        status = False
    else:
        status = True

    acc.disabled = status
    db.session.commit()
    return "OK"

@bp.route("/api/account/<int:acc_id>", methods=["DELETE"])
@login_required
@perm_manager.require('admin')
def api_account_delete(acc_id: int) -> Response:
    db.session.query(Account).filter(Account.id == acc_id).delete()
    db.session.commit()
    return flask.jsonify(status="OK")

@bp.route('/accounts/downloadlist/cvs')
@login_required
@perm_manager.require('accounts_download_list')
def accounts_download_csv() -> Response:
    def iter_accs(data):
        for account in data:
            for ci, char in enumerate(account.characters):
                if ci > 0:
                    yield ", " + char.eve_name
                else:
                    yield char.eve_name
            yield '\n'

    permission = perm_manager.get_permission('include_in_accountlist')
    include_check = (Account.disabled == False)
    role_check = None
    for role_need in permission.needs:
        if role_need.method != 'role':
            continue
        if role_check is None:
            role_check = (Role.name == role_need.value)
        else:
            role_check = role_check | (Role.name == role_need.value)

    include_check = (include_check & role_check)

    # noinspection PyPep8
    accs = db.session.query(Account).options(joinedload('characters')).join(Account.roles).filter(
        include_check ).order_by(
        Account.username).all()

    response = Response(iter_accs(accs), mimetype='text/csv')
    response.headers['Content-Disposition'] = 'attachment; filename=accounts.csv'
    return response

add_menu_entry('accounts.accounts', 'Accounts', perm_manager.get_permission('accounts_edit').can)
add_menu_entry('accounts.account_self', 'Own Settings', lambda: True)
<|endoftext|>


<commit_before># coding=utf-8
# Copyright 2020 The TensorFlow Datasets Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Methods to retrieve and store size/checksums associated to URLs.

"""

import os
from typing import Any, Dict, Iterable, List

import tensorflow.compat.v2 as tf

from tensorflow_datasets.core import utils

_ROOT_DIR = os.path.normpath(os.path.join(os.path.dirname(__file__), '../..'))

_CHECKSUM_DIRS = [
    os.path.join(_ROOT_DIR, 'url_checksums'),
]
_CHECKSUM_SUFFIX = '.txt'

class UrlInfo(object):  # TODO(tfds): Use dataclasses
  """Small wrapper around the url metadata (checksum, size).

  Attributes:
    size: Download size of the file
    checksum: Checksum of the file
  """

  def __init__(self, size: int, checksum: str):
    self.size = size
    self.checksum = checksum

  def asdict(self) -> Dict[str, Any]:
    """Returns the dict representation of the dataclass."""
    # TODO(tfds): Replace by `dataclasses.asdict(self)`
    return {
        'size': self.size,
        'checksum': self.checksum,
    }

  def __eq__(self, other) -> bool:
    return (
        type(self) == type(other) and  # pylint: disable=unidiomatic-typecheck
        self.size == other.size and
        self.checksum == other.checksum
    )

  def __ne__(self, other) -> bool:  # Required in Py2
    return not self == other

  def __repr__(self) -> str:
    return '{}(size={}, checksum={})'.format(
        type(self).__name__, self.size, self.checksum)

def add_checksums_dir(checksums_dir: str) -> None:
  """Registers a new checksums dir.

  This function allow external datasets not present in the tfds repository to
  define their own checksums_dir containing the dataset downloads checksums.

  Note: When redistributing your dataset, you should distribute the checksums
  files with it and set `add_checksums_dir` when the user is importing your
  `my_dataset.py`.

  ```
  # Set-up the folder containing the 'my_dataset.txt' checksums.
  checksum_dir = os.path.join(os.path.dirname(__file__), 'checksums/')
  checksum_dir = os.path.normpath(checksum_dir)

  # Add the checksum dir (will be executed when the user import your dataset)
  tfds.download.add_checksums_dir(checksum_dir)

  class MyDataset(tfds.core.DatasetBuilder):
    ...
  ```

  Args:
    checksums_dir: `str`, checksums dir to add to the registry
  """
  if checksums_dir in _CHECKSUM_DIRS:  # Avoid duplicate
    return
  _CHECKSUM_DIRS.append(checksums_dir)

def _list_dir(path: str) -> List[str]:
  return tf.io.gfile.listdir(path)

@utils.memoize()
def _checksum_paths() -> Dict[str, str]:
  """Returns dict {'dataset_name': 'path/to/checksums/file'}."""
  dataset2path = {}
  for dir_path in _CHECKSUM_DIRS:
    for fname in _list_dir(dir_path):
      if not fname.endswith(_CHECKSUM_SUFFIX):
        continue
      fpath = os.path.join(dir_path, fname)
      dataset_name = fname[:-len(_CHECKSUM_SUFFIX)]
      dataset2path[dataset_name] = fpath
  return dataset2path

def _get_path(dataset_name: str) -> str:
  """Returns path to where checksums are stored for a given dataset."""
  path = _checksum_paths().get(dataset_name, None)
  if path:
    return path
  msg = (
      'No checksums file could be find for dataset {}. Please '
      'create one in one of:\n{}'
      'If you are developing your own dataset outsite tfds, you can register '
      'your own checksums_dir with `tfds.download.add_checksums_dir('
      'checksums_dir)` or pass it to the download_and_prepare script with '
      '`--checksums_dir=`'
  ).format(
      dataset_name,
      ''.join(['* {}\n'.format(c) for c in _CHECKSUM_DIRS]))
  raise AssertionError(msg)

def _get_url_infos(checksums_path: str) -> Dict[str, UrlInfo]:
  """Returns {URL: (size, checksum)}s stored within file at given path."""
  with tf.io.gfile.GFile(checksums_path) as f:
    content = f.read()
  return _parse_url_infos(content.splitlines())

def _parse_url_infos(checksums_file: Iterable[str]) -> Dict[str, UrlInfo]:
  """Returns {URL: (size, checksum)}s stored within given file."""
  url_infos = {}
  for line in checksums_file:
    line = line.strip()  # Remove the trailing '\r' on Windows OS.
    if not line or line.startswith('#'):
      continue
    # URL might have spaces inside, but size and checksum will not.
    url, size, checksum = line.rsplit(' ', 2)
    url_infos[url] = UrlInfo(size=int(size), checksum=checksum)
  return url_infos

def url_infos_from_path(checksums_path: str) -> Dict[str, UrlInfo]:
  with tf.io.gfile.GFile(checksums_path) as f:
    return _parse_url_infos(f.read().splitlines())

@utils.memoize()
def get_all_url_infos() -> Dict[str, UrlInfo]:
  """Returns dict associating URL to (size, sha256)."""
  url_infos = {}
  for path in _checksum_paths().values():
    dataset_url_infos = _get_url_infos(path)
    for url, url_info in dataset_url_infos.items():
      if url_infos.get(url, url_info) != url_info:
        raise AssertionError(
            'URL {} is registered with 2+ distinct size/checksum tuples. '
            '{} vs {}'.format(url, url_info, url_infos[url]))
    url_infos.update(dataset_url_infos)
  return url_infos

def store_checksums(dataset_name: str, url_infos: Dict[str, UrlInfo]) -> None:
  """Store given checksums and sizes for specific dataset.

  Content of file is never disgarded, only updated. This is to ensure that if
  process is killed right after first download finishes, checksums registered
  during previous runs aren't lost.

  It is the responsibility of the caller not to call function multiple times in
  parallel for a given dataset.

  Only original file content is updated. This means the entire set of new sizes
  and checksums must be given at every call.

  Args:
    dataset_name: string.
    url_infos: dict, {url: (size_in_bytes, checksum)}.
  """
  path = _get_path(dataset_name)
  original_data = _get_url_infos(path)
  new_data = original_data.copy()
  new_data.update(url_infos)
  if original_data == new_data:
    return
  with tf.io.gfile.GFile(path, 'w') as f:
    for url, url_info in sorted(new_data.items()):
      f.write('{} {} {}\n'.format(url, url_info.size, url_info.checksum))
<commit_msg>Add core tab support for checksums<commit_after># coding=utf-8
# Copyright 2020 The TensorFlow Datasets Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Methods to retrieve and store size/checksums associated to URLs.

"""

import os
from typing import Any, Dict, Iterable, List

import tensorflow.compat.v2 as tf

from tensorflow_datasets.core import utils

_ROOT_DIR = os.path.normpath(os.path.join(os.path.dirname(__file__), '../..'))

_CHECKSUM_DIRS = [
    os.path.join(_ROOT_DIR, 'url_checksums'),
]
_CHECKSUM_SUFFIX = '.txt'

class UrlInfo(object):  # TODO(tfds): Use dataclasses
  """Small wrapper around the url metadata (checksum, size).

  Attributes:
    size: Download size of the file
    checksum: Checksum of the file
  """

  def __init__(self, size: int, checksum: str):
    self.size = size
    self.checksum = checksum

  def asdict(self) -> Dict[str, Any]:
    """Returns the dict representation of the dataclass."""
    # TODO(tfds): Replace by `dataclasses.asdict(self)`
    return {
        'size': self.size,
        'checksum': self.checksum,
    }

  def __eq__(self, other) -> bool:
    return (
        type(self) == type(other) and  # pylint: disable=unidiomatic-typecheck
        self.size == other.size and
        self.checksum == other.checksum
    )

  def __ne__(self, other) -> bool:  # Required in Py2
    return not self == other

  def __repr__(self) -> str:
    return '{}(size={}, checksum={})'.format(
        type(self).__name__, self.size, self.checksum)

def add_checksums_dir(checksums_dir: str) -> None:
  """Registers a new checksums dir.

  This function allow external datasets not present in the tfds repository to
  define their own checksums_dir containing the dataset downloads checksums.

  Note: When redistributing your dataset, you should distribute the checksums
  files with it and set `add_checksums_dir` when the user is importing your
  `my_dataset.py`.

  ```
  # Set-up the folder containing the 'my_dataset.txt' checksums.
  checksum_dir = os.path.join(os.path.dirname(__file__), 'checksums/')
  checksum_dir = os.path.normpath(checksum_dir)

  # Add the checksum dir (will be executed when the user import your dataset)
  tfds.download.add_checksums_dir(checksum_dir)

  class MyDataset(tfds.core.DatasetBuilder):
    ...
  ```

  Args:
    checksums_dir: `str`, checksums dir to add to the registry
  """
  if checksums_dir in _CHECKSUM_DIRS:  # Avoid duplicate
    return
  _CHECKSUM_DIRS.append(checksums_dir)

def _list_dir(path: str) -> List[str]:
  return tf.io.gfile.listdir(path)

@utils.memoize()
def _checksum_paths() -> Dict[str, str]:
  """Returns dict {'dataset_name': 'path/to/checksums/file'}."""
  dataset2path = {}
  for dir_path in _CHECKSUM_DIRS:
    for fname in _list_dir(dir_path):
      if not fname.endswith(_CHECKSUM_SUFFIX):
        continue
      fpath = os.path.join(dir_path, fname)
      dataset_name = fname[:-len(_CHECKSUM_SUFFIX)]
      dataset2path[dataset_name] = fpath
  return dataset2path

def _get_path(dataset_name: str) -> str:
  """Returns path to where checksums are stored for a given dataset."""
  path = _checksum_paths().get(dataset_name, None)
  if path:
    return path
  msg = (
      'No checksums file could be find for dataset {}. Please '
      'create one in one of:\n{}'
      'If you are developing your own dataset outsite tfds, you can register '
      'your own checksums_dir with `tfds.download.add_checksums_dir('
      'checksums_dir)` or pass it to the download_and_prepare script with '
      '`--checksums_dir=`'
  ).format(
      dataset_name,
      ''.join(['* {}\n'.format(c) for c in _CHECKSUM_DIRS]))
  raise AssertionError(msg)

def _get_url_infos(checksums_path: str) -> Dict[str, UrlInfo]:
  """Returns {URL: (size, checksum)}s stored within file at given path."""
  with tf.io.gfile.GFile(checksums_path) as f:
    content = f.read()
  return _parse_url_infos(content.splitlines())

def _parse_url_infos(checksums_file: Iterable[str]) -> Dict[str, UrlInfo]:
  """Returns {URL: (size, checksum)}s stored within given file."""
  url_infos = {}
  for line in checksums_file:
    line = line.strip()  # Remove the trailing '\r' on Windows OS.
    if not line or line.startswith('#'):
      continue
    # URL might have spaces inside, but size and checksum will not.
    url, size, checksum = line.rsplit('\t', 2)
    url_infos[url] = UrlInfo(size=int(size), checksum=checksum)
  return url_infos

@utils.memoize()
def get_all_url_infos() -> Dict[str, UrlInfo]:
  """Returns dict associating URL to (size, sha256)."""
  url_infos = {}
  for path in _checksum_paths().values():
    dataset_url_infos = _get_url_infos(path)
    for url, url_info in dataset_url_infos.items():
      if url_infos.get(url, url_info) != url_info:
        raise AssertionError(
            'URL {} is registered with 2+ distinct size/checksum tuples. '
            '{} vs {}'.format(url, url_info, url_infos[url]))
    url_infos.update(dataset_url_infos)
  return url_infos

def store_checksums(dataset_name: str, url_infos: Dict[str, UrlInfo]) -> None:
  """Store given checksums and sizes for specific dataset.

  Content of file is never disgarded, only updated. This is to ensure that if
  process is killed right after first download finishes, checksums registered
  during previous runs aren't lost.

  It is the responsibility of the caller not to call function multiple times in
  parallel for a given dataset.

  Only original file content is updated. This means the entire set of new sizes
  and checksums must be given at every call.

  Args:
    dataset_name: string.
    url_infos: dict, {url: (size_in_bytes, checksum)}.
  """
  path = _get_path(dataset_name)
  original_data = _get_url_infos(path)
  new_data = original_data.copy()
  new_data.update(url_infos)
  if original_data == new_data:
    return
  with tf.io.gfile.GFile(path, 'w') as f:
    for url, url_info in sorted(new_data.items()):
      f.write('{}\t{}\t{}\n'.format(url, url_info.size, url_info.checksum))
<|endoftext|>


<commit_before># vim: set expandtab sw=4 ts=4:
# disable pylint message about unused variable 'fig'
# pylint: disable=unused-variable
'''
Generates a trend (graph) from the buildtimes in buildtimes.xml

Copyright (C) 2014 Dieter Adriaenssens <ruleant@users.sourceforge.net>

This file is part of buildtime-trend
<https://github.com/ruleant/buildtime-trend/>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
'''

from lxml import etree
import matplotlib
# Force matplotlib to not use any Xwindow backend.
matplotlib.use('Agg')
from matplotlib import pyplot as plt
from buildtimetrend.tools import check_file

class Trend(object):
    '''
    Trend class, generates a chart from gathered buildtime data
    '''
    def __init__(self):
        self.stages = {}
        self.builds = []

    def gather_data(self, result_file):
        '''
        Get buildtime data from an xml file

        Parameters
        - result_file : xml file containing the buildtime data
        '''
        # load buildtimes file
        if check_file(result_file):
            root_xml = etree.parse(result_file).getroot()
        else:
            return False

        index = 0
        # print content of buildtimes file
        for build_xml in root_xml:
            build_id = build_xml.get('build')
            job_id = build_xml.get('job')

            if job_id is None and build_id is None:
                build_id = "#%d" % (index + 1)
            elif job_id is not None:
                build_id = job_id

            build_summary = "Build ID : %s, Job : %s" % (build_id, job_id)

            self.builds.append(build_id)

            # add 0 to each existing stage, to make sure that
            # the indexes of each value
            # are correct, even if a stage does not exist in a build
            # if a stage exists, the zero will be replaced by its duration
            for stage in self.stages:
                self.stages[stage].append(0)

            # add duration of each stage to stages list
            for build_child in build_xml:
                if build_child.tag == 'stages':
                    build_summary += ", stages : " + str(len(build_child))
                    for stage in build_child:
                        if (stage.tag == 'stage' and
                                stage.get('name') is not None and
                                stage.get('duration') is not None):
                            if stage.get('name') in self.stages:
                                temp_dict = self.stages[stage.get('name')]
                            else:
                                # when a new stage is added,
                                # create list with zeros,
                                # one for each existing build
                                temp_dict = [0]*(index + 1)
                            temp_dict[index] = int(stage.get('duration'))
                            self.stages[stage.get('name')] = temp_dict
            print build_summary
            index += 1
        return True

    def generate(self, trend_file):
        '''
        Generates the trend chart and saves it as a PNG image using matplotlib

        Parameters
        - trend_file : file name to save chart image to
        '''
        fig, axes = plt.subplots()

        # add data
        x_values = range(len(self.builds))
        plots = plt.stackplot(x_values, self.stages.values())
        plt.xticks(x_values, self.builds, rotation=45, size=10)

        # label axes and add graph title
        axes.set_xlabel("Builds", {'fontsize': 14})
        axes.xaxis.set_label_coords(1.05, -0.05)
        axes.set_ylabel("Duration [s]", {'fontsize': 14})
        axes.set_title("Build stages trend", {'fontsize': 22})

        # display legend
        legend_proxies = []
        for plot in plots:
            legend_proxies.append(
                plt.Rectangle((0, 0), 1, 1, fc=plot.get_facecolor()[0]))
        # add legend in reverse order, in upper left corner
        axes.legend(legend_proxies[::-1], self.stages.keys()[::-1], loc=2)

        # save figure
        plt.savefig(trend_file)
<commit_msg>build_name is not the same as build_id<commit_after># vim: set expandtab sw=4 ts=4:
# disable pylint message about unused variable 'fig'
# pylint: disable=unused-variable
'''
Generates a trend (graph) from the buildtimes in buildtimes.xml

Copyright (C) 2014 Dieter Adriaenssens <ruleant@users.sourceforge.net>

This file is part of buildtime-trend
<https://github.com/ruleant/buildtime-trend/>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
'''

from lxml import etree
import matplotlib
# Force matplotlib to not use any Xwindow backend.
matplotlib.use('Agg')
from matplotlib import pyplot as plt
from buildtimetrend.tools import check_file

class Trend(object):
    '''
    Trend class, generates a chart from gathered buildtime data
    '''
    def __init__(self):
        self.stages = {}
        self.builds = []

    def gather_data(self, result_file):
        '''
        Get buildtime data from an xml file

        Parameters
        - result_file : xml file containing the buildtime data
        '''
        # load buildtimes file
        if check_file(result_file):
            root_xml = etree.parse(result_file).getroot()
        else:
            return False

        index = 0
        # print content of buildtimes file
        for build_xml in root_xml:
            build_id = build_xml.get('build')
            job_id = build_xml.get('job')

            if job_id is None and build_id is None:
                build_name = "#%d" % (index + 1)
            elif job_id is not None:
                build_name = job_id
            else:
                build_name = build_id

            build_summary = "Build ID : %s, Job : %s" % (build_id, job_id)

            self.builds.append(build_name)

            # add 0 to each existing stage, to make sure that
            # the indexes of each value
            # are correct, even if a stage does not exist in a build
            # if a stage exists, the zero will be replaced by its duration
            for stage in self.stages:
                self.stages[stage].append(0)

            # add duration of each stage to stages list
            for build_child in build_xml:
                if build_child.tag == 'stages':
                    build_summary += ", stages : " + str(len(build_child))
                    for stage in build_child:
                        if (stage.tag == 'stage' and
                                stage.get('name') is not None and
                                stage.get('duration') is not None):
                            if stage.get('name') in self.stages:
                                temp_dict = self.stages[stage.get('name')]
                            else:
                                # when a new stage is added,
                                # create list with zeros,
                                # one for each existing build
                                temp_dict = [0]*(index + 1)
                            temp_dict[index] = int(stage.get('duration'))
                            self.stages[stage.get('name')] = temp_dict
            print build_summary
            index += 1
        return True

    def generate(self, trend_file):
        '''
        Generates the trend chart and saves it as a PNG image using matplotlib

        Parameters
        - trend_file : file name to save chart image to
        '''
        fig, axes = plt.subplots()

        # add data
        x_values = range(len(self.builds))
        plots = plt.stackplot(x_values, self.stages.values())
        plt.xticks(x_values, self.builds, rotation=45, size=10)

        # label axes and add graph title
        axes.set_xlabel("Builds", {'fontsize': 14})
        axes.xaxis.set_label_coords(1.05, -0.05)
        axes.set_ylabel("Duration [s]", {'fontsize': 14})
        axes.set_title("Build stages trend", {'fontsize': 22})

        # display legend
        legend_proxies = []
        for plot in plots:
            legend_proxies.append(
                plt.Rectangle((0, 0), 1, 1, fc=plot.get_facecolor()[0]))
        # add legend in reverse order, in upper left corner
        axes.legend(legend_proxies[::-1], self.stages.keys()[::-1], loc=2)

        # save figure
        plt.savefig(trend_file)
<|endoftext|>