<commit_before>#!/usr/bin/env python3
import argparse
# import glob
# import math
# import re
import os
from collections import namedtuple
from functools import lru_cache
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
# import matplotlib.ticker as ticker
import pandas as pd

import artistools as at

DEFAULTSPECPATH = '../example_run/spec.out'
defaultoutputfile = 'plotnonthermal_cell{0:03d}_timestep{1:03d}.pdf'

@lru_cache(maxsize=4)
def read_files(modelpath, timestep=-1, modelgridindex=-1):
    """Read ARTIS -thermal spectrum data into a pandas DataFrame."""
    nonthermaldata = pd.DataFrame()

    mpiranklist = at.get_mpiranklist(modelpath, modelgridindex=modelgridindex)
    for folderpath in at.get_runfolders(modelpath, timestep=timestep):
        for mpirank in mpiranklist:
            nonthermalfile = f'nonthermalspec_{mpirank:04d}.out'
            filepath = Path(folderpath, nonthermalfile)
            if not filepath.is_file():
                filepath = Path(folderpath, nonthermalfile + '.gz')
                if not filepath.is_file():
                    print(f'Warning: Could not find {filepath.relative_to(modelpath.parent)}')
                    continue

            if modelgridindex > -1:
                filesize = Path(filepath).stat().st_size / 1024 / 1024
                print(f'Reading {Path(filepath).relative_to(modelpath.parent)} ({filesize:.2f} MiB)')

            nonthermaldata_thisfile = pd.read_csv(filepath, delim_whitespace=True, error_bad_lines=False)
            # radfielddata_thisfile[['modelgridindex', 'timestep']].apply(pd.to_numeric)

            if timestep >= 0:
                nonthermaldata_thisfile.query('timestep==@timestep', inplace=True)

            if modelgridindex >= 0:
                nonthermaldata_thisfile.query('modelgridindex==@modelgridindex', inplace=True)

            if not nonthermaldata_thisfile.empty:
                if timestep >= 0 and modelgridindex >= 0:
                    return nonthermaldata_thisfile
                else:
                    nonthermaldata = nonthermaldata.append(nonthermaldata_thisfile.copy(), ignore_index=True)

    return nonthermaldata

def ar_xs(energy_ev, ionpot_ev, A, B, C, D):
    u = energy_ev / ionpot_ev
    if u <= 1:
        return 0
    return 1e-14 * (
        A * (1 - 1 / u) + B * pow((1 - 1 / u), 2) + C * np.log(u) + D * np.log(u) / u) / (u * pow(ionpot_ev, 2))

def xs_fe2_old(energy):
    # AR1985
    shell_a = ar_xs(energy, 16.2, 90.0, -60.0, 0.2, -86)
    shell_b = ar_xs(energy, 17.5, 18.6, -5.9, 0.6, -9)
    shell_c = ar_xs(energy, 81, 69.9, -23.7, 9.5, -51.7)
    return shell_a + shell_b + shell_c

def get_arxs_array_shell(arr_enev, row):
    ar_xs_array = np.zeros(len(arr_enev))

    ionpot_ev, A, B, C, D = row.ionpot_ev, row.A, row.B, row.C, row.D
    for index, energy_ev in enumerate(arr_enev):
        u = energy_ev / ionpot_ev
        if u <= 1:
            ar_xs_array[index] = 0.
        else:
            ar_xs_array[index] = (1e-14 * (A * (1 - 1 / u) + B * pow((1 - 1 / u), 2) +
                                           C * np.log(u) + D * np.log(u) / u) / (u * pow(ionpot_ev, 2)))

    return ar_xs_array

def get_arxs_array(arr_enev, dfcollion, Z, ionstage):
    ar_xs_array = np.zeros(len(arr_enev))
    dfcollion_thision = dfcollion.query('Z == @Z and ionstage == @ionstage')
    print(dfcollion_thision)
    for index, row in dfcollion_thision.iterrows():
        ar_xs_array = np.add(ar_xs_array, np.array(
            [ar_xs(energy_ev, row.ionpot_ev, row.A, row.B, row.C, row.D) for energy_ev in arr_enev]))
    return ar_xs_array

def read_colliondata(collionfilename='collion.txt'):
    collionrow = namedtuple('collionrow', ['Z', 'nelec', 'n', 'l', 'ionpot_ev', 'A', 'B', 'C', 'D'])
    with open(os.path.join(at.PYDIR, 'data', collionfilename), 'r') as collionfile:
        print(f'Collionfile: expecting {collionfile.readline().strip()} rows')
        dfcollion = pd.read_csv(
            collionfile, delim_whitespace=True, header=None, names=collionrow._fields)
    dfcollion.eval('ionstage = Z - nelec + 1', inplace=True)

    return dfcollion

def make_xs_plot(axis, nonthermaldata, timestep, outputfile, args):
    dfcollion = read_colliondata()

    arr_en = nonthermaldata['energy_ev'].unique()

    # arr_xs_old = [xs_fe2_old(en) for en in arr_en]
    # arr_xs_times_y = [xs_fe1(en) * y for en, y in zip(nonthermaldata['energy_ev'], nonthermaldata['y'])]

    axis.plot(arr_en, get_arxs_array(arr_en, dfcollion, 26, 2), linewidth=2.0, label='Fe II')
    axis.plot(arr_en, get_arxs_array(arr_en, dfcollion, 28, 2), linewidth=2.0, label='Ni II')

    axis.set_ylabel(r'cross section (cm2)')

    if not args.nolegend:
        axis.legend(loc='upper center', handlelength=2, frameon=False, numpoints=1, prop={'size': 13})

def make_plot(modelpaths, args):
              # nonthermaldata, modelpath, modelgridindex, timestep, outputfile, args):
    """Draw the bin edges, fitted field, and emergent spectrum."""
    nplots = 1 if not args.xsplot else 2
    fig, axes = plt.subplots(nrows=nplots, ncols=1, sharex=True,
                             figsize=(args.figscale * at.figwidth, args.figscale * at.figwidth * 0.7 * nplots),
                             tight_layout={"pad": 0.2, "w_pad": 0.0, "h_pad": 0.0})

    if nplots == 1:
        axes = [axes]

    if args.kf1992spec:
        kf92spec = pd.read_csv(
            Path(modelpaths[0], 'KF1992spec-fig1.txt'),
            header=None, names=['e_kev', 'log10_y'])
        kf92spec['energy_ev'] = kf92spec['e_kev'] * 1000.
        kf92spec.eval('y = 10 ** log10_y', inplace=True)
        axes[0].plot(kf92spec['energy_ev'], kf92spec['log10_y'],
                     linewidth=2.0, color='red', label='Kozma & Fransson (1992)')

    for index, modelpath in enumerate(modelpaths):
        modelname = at.get_model_name(modelpath)
        if args.velocity >= 0.:
            modelgridindex = at.get_mgi_of_velocity_kms(modelpath, args.velocity)
        else:
            modelgridindex = args.modelgridindex

        if args.timedays:
            timestep = at.get_timestep_of_timedays(modelpath, args.timedays)
        else:
            timestep = args.timestep

        nonthermaldata = read_files(
            modelpath=Path(modelpath),
            modelgridindex=modelgridindex, timestep=timestep)

        if nonthermaldata.empty:
            print(f'No data for timestep {timestep:d}')
            continue

        if index < len(args.modellabels):
            model_label = args.modellabels[index]
        else:
            model_label = f'{modelname} cell {modelgridindex} at timestep {timestep}'
            try:
                time_days = float(at.get_timestep_time('.', timestep))
            except FileNotFoundError:
                time_days = 0
            else:
                model_label += f' ({time_days:.2f}d)'

        outputfile = str(args.outputfile).format(modelgridindex, timestep)
        print(f'Plotting timestep {timestep:d}')
        # ymax = max(nonthermaldata['y'])

        # nonthermaldata.plot(x='energy_ev', y='y', linewidth=1.5, ax=axis, color='blue', legend=False)
        axes[0].plot((nonthermaldata['energy_ev']), np.log10(nonthermaldata['y']), label=model_label,
                     linewidth=2.0, color='black' if index == 0 else None, alpha=0.95)
        axes[0].set_ylabel(r'log [y (e$^-$ / cm$^2$ / s / eV)]')

        if args.xsplot:
            make_xs_plot(axes[1], nonthermaldata, timestep, outputfile, args)

    if not args.nolegend:
        axes[0].legend(loc='best', handlelength=2, frameon=False, numpoints=1)

    axes[-1].set_xlabel(r'Energy (eV)')
    # axis.yaxis.set_minor_locator(ticker.MultipleLocator(base=0.1))
    # axis.set_yscale("log", nonposy='clip')
    for ax in axes:
        ax.tick_params(which='both', direction='in')
        if args.xmin is not None:
            ax.set_xlim(left=args.xmin)
        if args.xmax:
            ax.set_xlim(right=args.xmax)
    # axis.set_ylim(bottom=0.0, top=ymax)

    # axis.legend(loc='upper center', handlelength=2,
    #             frameon=False, numpoints=1, prop={'size': 13})

    print(f'Saving to {outputfile:s}')
    fig.savefig(outputfile, format='pdf')
    plt.close()

def addargs(parser):
    parser.add_argument('-modelpath', default=[], nargs='*', action=at.AppendPath,
                        help='Paths to ARTIS folders with spec.out or packets files')

    parser.add_argument('-modellabels', default=[], nargs='*',
                        help='Model name overrides')

    parser.add_argument('-listtimesteps', action='store_true',
                        help='Show the times at each timestep')

    parser.add_argument('-xsplot', action='store_true',
                        help='Show the cross section plot')

    parser.add_argument('-timedays', '-time', '-t',
                        help='Time in days to plot')

    parser.add_argument('-timestep', '-ts', type=int, default=-1,
                        help='Timestep number to plot')

    parser.add_argument('-modelgridindex', '-cell', type=int, default=0,
                        help='Modelgridindex to plot')

    parser.add_argument('-velocity', '-v', type=float, default=-1,
                        help='Specify cell by velocity')

    parser.add_argument('-xmin', type=float, default=0.,
                        help='Plot range: minimum energy in eV')

    parser.add_argument('-xmax', type=float,
                        help='Plot range: maximum energy in eV')

    parser.add_argument('--nolegend', action='store_true',
                        help='Suppress the legend from the plot')

    parser.add_argument('--kf1992spec', action='store_true',
                        help='Show the pure-oxygen result form Figure 1 of Kozma & Fransson 1992')

    parser.add_argument('-figscale', type=float, default=1.,
                        help='Scale factor for plot area. 1.0 is for single-column')

    parser.add_argument('-o', action='store', dest='outputfile', type=Path,
                        default=defaultoutputfile,
                        help='Filename for PDF file')

def main(args=None, argsraw=None, **kwargs):
    """Plot the electron energy distribution."""
    if args is None:
        parser = argparse.ArgumentParser(
            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
            description='Plot ARTIS non-thermal electron energy spectrum.')

        addargs(parser)
        parser.set_defaults(**kwargs)
        args = parser.parse_args(argsraw)

    if not args.modelpath:
        args.modelpath = [Path('.')]
    elif isinstance(args.modelpath, (str, Path)):
        args.modelpath = [args.modelpath]

    # flatten the list
    modelpaths = []
    for elem in args.modelpath:
        if isinstance(elem, list):
            modelpaths.extend(elem)
        else:
            modelpaths.append(elem)

    if os.path.isdir(args.outputfile):
        args.outputfile = os.path.join(args.outputfile, defaultoutputfile)

    if args.listtimesteps:
        at.showtimesteptimes()
    else:
        make_plot(modelpaths, args)

if __name__ == "__main__":
    main()
<commit_msg>Add non-thermal contribution plot option<commit_after>#!/usr/bin/env python3
import argparse
# import glob
# import math
# import re
import os
from collections import namedtuple
from functools import lru_cache
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
# import matplotlib.ticker as ticker
import pandas as pd
from astropy import units as u

import artistools as at
import artistools.spencerfano

DEFAULTSPECPATH = '../example_run/spec.out'
defaultoutputfile = 'plotnonthermal_cell{0:03d}_timestep{1:03d}.pdf'

@lru_cache(maxsize=4)
def read_files(modelpath, timestep=-1, modelgridindex=-1):
    """Read ARTIS -thermal spectrum data into a pandas DataFrame."""
    nonthermaldata = pd.DataFrame()

    mpiranklist = at.get_mpiranklist(modelpath, modelgridindex=modelgridindex)
    for folderpath in at.get_runfolders(modelpath, timestep=timestep):
        for mpirank in mpiranklist:
            nonthermalfile = f'nonthermalspec_{mpirank:04d}.out'
            filepath = Path(folderpath, nonthermalfile)
            if not filepath.is_file():
                filepath = Path(folderpath, nonthermalfile + '.gz')
                if not filepath.is_file():
                    print(f'Warning: Could not find {filepath.relative_to(modelpath.parent)}')
                    continue

            if modelgridindex > -1:
                filesize = Path(filepath).stat().st_size / 1024 / 1024
                print(f'Reading {Path(filepath).relative_to(modelpath.parent)} ({filesize:.2f} MiB)')

            nonthermaldata_thisfile = pd.read_csv(filepath, delim_whitespace=True, error_bad_lines=False)
            # radfielddata_thisfile[['modelgridindex', 'timestep']].apply(pd.to_numeric)

            if timestep >= 0:
                nonthermaldata_thisfile.query('timestep==@timestep', inplace=True)

            if modelgridindex >= 0:
                nonthermaldata_thisfile.query('modelgridindex==@modelgridindex', inplace=True)

            if not nonthermaldata_thisfile.empty:
                if timestep >= 0 and modelgridindex >= 0:
                    return nonthermaldata_thisfile
                else:
                    nonthermaldata = nonthermaldata.append(nonthermaldata_thisfile.copy(), ignore_index=True)

    return nonthermaldata

def ar_xs(energy_ev, ionpot_ev, A, B, C, D):
    u = energy_ev / ionpot_ev
    if u <= 1:
        return 0
    return 1e-14 * (
        A * (1 - 1 / u) + B * pow((1 - 1 / u), 2) + C * np.log(u) + D * np.log(u) / u) / (u * pow(ionpot_ev, 2))

def xs_fe2_old(energy):
    # AR1985
    shell_a = ar_xs(energy, 16.2, 90.0, -60.0, 0.2, -86)
    shell_b = ar_xs(energy, 17.5, 18.6, -5.9, 0.6, -9)
    shell_c = ar_xs(energy, 81, 69.9, -23.7, 9.5, -51.7)
    return shell_a + shell_b + shell_c

def get_arxs_array_shell(arr_enev, row):
    ar_xs_array = np.zeros(len(arr_enev))

    ionpot_ev, A, B, C, D = row.ionpot_ev, row.A, row.B, row.C, row.D
    for index, energy_ev in enumerate(arr_enev):
        u = energy_ev / ionpot_ev
        if u <= 1:
            ar_xs_array[index] = 0.
        else:
            ar_xs_array[index] = (1e-14 * (A * (1 - 1 / u) + B * pow((1 - 1 / u), 2) +
                                           C * np.log(u) + D * np.log(u) / u) / (u * pow(ionpot_ev, 2)))

    return ar_xs_array

def get_arxs_array(arr_enev, dfcollion, Z, ionstage):
    ar_xs_array = np.zeros(len(arr_enev))
    dfcollion_thision = dfcollion.query('Z == @Z and ionstage == @ionstage')
    print(dfcollion_thision)
    for index, row in dfcollion_thision.iterrows():
        ar_xs_array = np.add(ar_xs_array, np.array(
            [ar_xs(energy_ev, row.ionpot_ev, row.A, row.B, row.C, row.D) for energy_ev in arr_enev]))
    return ar_xs_array

def read_colliondata(collionfilename='collion.txt'):
    collionrow = namedtuple('collionrow', ['Z', 'nelec', 'n', 'l', 'ionpot_ev', 'A', 'B', 'C', 'D'])
    with open(os.path.join(at.PYDIR, 'data', collionfilename), 'r') as collionfile:
        print(f'Collionfile: expecting {collionfile.readline().strip()} rows')
        dfcollion = pd.read_csv(
            collionfile, delim_whitespace=True, header=None, names=collionrow._fields)
    dfcollion.eval('ionstage = Z - nelec + 1', inplace=True)

    return dfcollion

def make_xs_plot(axis, nonthermaldata, args):
    dfcollion = read_colliondata()

    arr_en = nonthermaldata['energy_ev'].unique()

    # arr_xs_old = [xs_fe2_old(en) for en in arr_en]
    # arr_xs_times_y = [xs_fe1(en) * y for en, y in zip(nonthermaldata['energy_ev'], nonthermaldata['y'])]

    axis.plot(arr_en, get_arxs_array(arr_en, dfcollion, 26, 2), linewidth=2.0, label='Fe II')
    axis.plot(arr_en, get_arxs_array(arr_en, dfcollion, 28, 2), linewidth=2.0, label='Ni II')

    axis.set_ylabel(r'cross section (cm2)')

    if not args.nolegend:
        axis.legend(loc='upper center', handlelength=2, frameon=False, numpoints=1, prop={'size': 13})

def inteuler(x, y):
    dx = y[1:] - y[:-1]
    return np.dot(x[:-1], dx)

def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata, args):
    estimators = at.estimators.read_estimators(modelpath, get_ion_values=True, get_heatingcooling=True,
                                               modelgridindex=modelgridindex, timestep=timestep)

    # print(estimators[(timestep, modelgridindex)].keys())
    total_depev = (estimators[(timestep, modelgridindex)]['total_dep'] * u.erg.to('eV'))

    print(f"Deposition: {total_depev:.1f} eV / cm3 / s")

    arr_enev = nonthermaldata['energy_ev'].values
    arr_y = nonthermaldata['y'].values

    frac_ionisation = 0.

    dfcollion = read_colliondata()

    elementlist = at.get_composition_data(modelpath)
    totalpop = estimators[(timestep, modelgridindex)]['populations']['total']
    nelements = len(elementlist)
    for element in range(nelements):
        Z = elementlist.Z[element]

        elpop = estimators[(timestep, modelgridindex)]['populations'][Z]
        if elpop <= 1e-4 * totalpop:
            continue

        arr_ionisation_element = np.zeros(len(arr_enev), dtype=np.float)
        frac_ionisation_element = 0.

        nions = elementlist.nions[element]
        for ion in range(nions):
            ionstage = ion + elementlist.lowermost_ionstage[element]
            ionpop = estimators[(timestep, modelgridindex)]['populations'][(Z, ionstage)]

            dfcollion_thision = dfcollion.query('Z == @Z and ionstage == @ionstage')

            # print(at.get_ionstring(Z, ionstage), ionpop)

            arr_ionisation_ion = np.zeros(len(arr_enev), dtype=np.float)
            frac_ionisation_ion = 0.

            for index, row in dfcollion_thision.iterrows():
                arr_xs = get_arxs_array_shell(arr_enev, row)
                arr_ionisation_shell = ionpop * arr_y * arr_xs * row.ionpot_ev / total_depev
                arr_ionisation_ion += arr_ionisation_shell

                frac_ionisation_shell = np.trapz(x=arr_enev, y=arr_ionisation_shell)
                frac_ionisation_ion += frac_ionisation_shell

            arr_ionisation_element += arr_ionisation_ion
            frac_ionisation_element += frac_ionisation_ion

        frac_ionisation += frac_ionisation_element

        if frac_ionisation_element > 1e-5:
            axis.plot(arr_enev, arr_ionisation_element, label=f'Ionisation Z={Z}')

    nne = estimators[(timestep, modelgridindex)]['nne']
    arr_heating = np.array([at.spencerfano.lossfunction(enev, nne) / total_depev for enev in arr_enev])

    frac_heating = np.trapz(x=arr_enev, y=arr_heating)

    print(f'   frac_heating: {frac_heating}')
    print(f'frac_ionisation: {frac_ionisation}')

    axis.plot(arr_enev, arr_heating, label=f'Heating')

    axis.legend(loc='best', handlelength=2, frameon=False, numpoints=1, prop={'size': 11})

def make_plot(modelpaths, args):
    nplots = 1
    if args.xsplot:
        nplots += 1
    if args.showcontributions:
        nplots += 1
    fig, axes = plt.subplots(nrows=nplots, ncols=1, sharex=True,
                             figsize=(args.figscale * at.figwidth, args.figscale * at.figwidth * 0.7 * nplots),
                             tight_layout={"pad": 0.2, "w_pad": 0.0, "h_pad": 0.0})

    if nplots == 1:
        axes = [axes]

    if args.kf1992spec:
        kf92spec = pd.read_csv(Path(modelpaths[0], 'KF1992spec-fig1.txt'), header=None, names=['e_kev', 'log10_y'])
        kf92spec['energy_ev'] = kf92spec['e_kev'] * 1000.
        kf92spec.eval('y = 10 ** log10_y', inplace=True)
        axes[0].plot(kf92spec['energy_ev'], kf92spec['log10_y'],
                     linewidth=2.0, color='red', label='Kozma & Fransson (1992)')

    for index, modelpath in enumerate(modelpaths):
        modelname = at.get_model_name(modelpath)
        if args.velocity >= 0.:
            modelgridindex = at.get_mgi_of_velocity_kms(modelpath, args.velocity)
        else:
            modelgridindex = args.modelgridindex

        if args.timedays:
            timestep = at.get_timestep_of_timedays(modelpath, args.timedays)
        else:
            timestep = args.timestep

        nonthermaldata = read_files(
            modelpath=Path(modelpath),
            modelgridindex=modelgridindex, timestep=timestep)

        if args.xmin:
            nonthermaldata.query('energy_ev >= @args.xmin', inplace=True)

        if nonthermaldata.empty:
            print(f'No data for timestep {timestep:d}')
            continue

        if index < len(args.modellabels):
            model_label = args.modellabels[index]
        else:
            model_label = f'{modelname} cell {modelgridindex} at timestep {timestep}'
            try:
                time_days = float(at.get_timestep_time('.', timestep))
            except FileNotFoundError:
                time_days = 0
            else:
                model_label += f' ({time_days:.2f}d)'

        outputfile = str(args.outputfile).format(modelgridindex, timestep)
        print(f'Plotting timestep {timestep:d}')
        # ymax = max(nonthermaldata['y'])

        # nonthermaldata.plot(x='energy_ev', y='y', linewidth=1.5, ax=axis, color='blue', legend=False)
        axes[0].plot((nonthermaldata['energy_ev']), np.log10(nonthermaldata['y']), label=model_label,
                     linewidth=2.0, color='black' if index == 0 else None, alpha=0.95)
        axes[0].set_ylabel(r'log [y (e$^-$ / cm$^2$ / s / eV)]')

        if args.showcontributions:
            plot_contributions(axes[1], modelpath, timestep, modelgridindex, nonthermaldata, args)

        if args.xsplot:
            make_xs_plot(axes[-1], nonthermaldata, args)

    if not args.nolegend:
        axes[0].legend(loc='best', handlelength=2, frameon=False, numpoints=1)

    axes[-1].set_xlabel(r'Energy (eV)')
    # axis.yaxis.set_minor_locator(ticker.MultipleLocator(base=0.1))
    # axis.set_yscale("log", nonposy='clip')
    for ax in axes:
        ax.tick_params(which='both', direction='in')
        if args.xmin is not None:
            ax.set_xlim(left=args.xmin)
        if args.xmax:
            ax.set_xlim(right=args.xmax)
    # axis.set_ylim(bottom=0.0, top=ymax)

    # axis.legend(loc='upper center', handlelength=2,
    #             frameon=False, numpoints=1, prop={'size': 13})

    print(f'Saving to {outputfile:s}')
    fig.savefig(outputfile, format='pdf')
    plt.close()

def addargs(parser):
    parser.add_argument('-modelpath', default=[], nargs='*', action=at.AppendPath,
                        help='Paths to ARTIS folders with spec.out or packets files')

    parser.add_argument('-modellabels', default=[], nargs='*',
                        help='Model name overrides')

    parser.add_argument('-listtimesteps', action='store_true',
                        help='Show the times at each timestep')

    parser.add_argument('-xsplot', action='store_true',
                        help='Show the cross section plot')

    parser.add_argument('-timedays', '-time', '-t',
                        help='Time in days to plot')

    parser.add_argument('-timestep', '-ts', type=int, default=-1,
                        help='Timestep number to plot')

    parser.add_argument('-modelgridindex', '-cell', type=int, default=0,
                        help='Modelgridindex to plot')

    parser.add_argument('-velocity', '-v', type=float, default=-1,
                        help='Specify cell by velocity')

    parser.add_argument('-xmin', type=float, default=0.,
                        help='Plot range: minimum energy in eV')

    parser.add_argument('-xmax', type=float,
                        help='Plot range: maximum energy in eV')

    parser.add_argument('--nolegend', action='store_true',
                        help='Suppress the legend from the plot')

    parser.add_argument('--showcontributions', action='store_true',
                        help='Plot the NT contributions to ionisation and heating energy')

    parser.add_argument('--kf1992spec', action='store_true',
                        help='Show the pure-oxygen result form Figure 1 of Kozma & Fransson 1992')

    parser.add_argument('-figscale', type=float, default=1.,
                        help='Scale factor for plot area. 1.0 is for single-column')

    parser.add_argument('-o', action='store', dest='outputfile', type=Path,
                        default=defaultoutputfile,
                        help='Filename for PDF file')

def main(args=None, argsraw=None, **kwargs):
    """Plot the electron energy distribution."""
    if args is None:
        parser = argparse.ArgumentParser(
            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
            description='Plot ARTIS non-thermal electron energy spectrum.')

        addargs(parser)
        parser.set_defaults(**kwargs)
        args = parser.parse_args(argsraw)

    if not args.modelpath:
        args.modelpath = [Path('.')]
    elif isinstance(args.modelpath, (str, Path)):
        args.modelpath = [args.modelpath]

    # flatten the list
    modelpaths = []
    for elem in args.modelpath:
        if isinstance(elem, list):
            modelpaths.extend(elem)
        else:
            modelpaths.append(elem)

    if os.path.isdir(args.outputfile):
        args.outputfile = os.path.join(args.outputfile, defaultoutputfile)

    if args.listtimesteps:
        at.showtimesteptimes()
    else:
        make_plot(modelpaths, args)

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


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import os
from collections import namedtuple
from Cookie import SimpleCookie
from StringIO import StringIO

from aspen.network_engines import ThreadedBuffer
from aspen.http.request import Request
from aspen.sockets.channel import Channel
from aspen.sockets.socket import Socket
from aspen.sockets.transport import XHRPollingTransport
from aspen.website import Website

BOUNDARY = 'BoUnDaRyStRiNg'
MULTIPART_CONTENT = 'multipart/form-data; boundary=%s' % BOUNDARY

def encode_multipart(boundary, data):
    """
    Encodes multipart POST data from a dictionary of form values.

    Borrowed from Django
    The key will be used as the form data name; the value will be transmitted
    as content. If the value is a file, the contents of the file will be sent
    as an application/octet-stream; otherwise, str(value) will be sent.
    """
    lines = []

    for (key, value) in data.items():
        lines.extend([
            '--' + boundary,
            'Content-Disposition: form-data; name="%s"' % str(key),
            '',
            str(value)
        ])

    lines.extend([
        '--' + boundary + '--',
        '',
    ])
    return '\r\n'.join(lines)

class Harness(object):
    """
    The Aspen testing harness.

    Used in tests to emulate ``GET`` and ``POST`` requests by sending them
    into a ``Website`` instance's ``respond`` method.

    Aspen does not define any User data structures or modules. If you want to
    do anything with users/sessions etc in your tests it is expected that you
    will subclass this class and add a ``add_cookie_info`` method.

    For example, in gittip a suitable subclass might be::

        class GittipTestClient(TestClient):

            def add_cookie_info(self, request, cookie_info):
                if cookie_info:
                    user = cookie_info.get('user')
                    if user is not None:
                        user = User.from_id(user)
                        # Note that Cookie needs a bytestring.
                        request.headers.cookie['session'] = user.session_token

    Example usage in a test::

        def test_api_handles_posts():
            client = TestClient(website)

            # We need to get ourselves a token!
            response = client.get('/')
            csrf_token = response.request.context['csrf_token']

            # Then, add a $1.50 and $3.00 tip
            response = client.post("/test_tippee1/tip.json",
                                {'amount': "1.00", 'csrf_token': csrf_token},
                                cookie_info={'user': 'test_tipper'})

            # Confirm we get back the right amounts in the JSON body.
            first_data = json.loads(response.body)
            assert_equal(first_data['amount'], "1.00")
    """

    def __init__(self, www, project):
        self.fs = namedtuple('fs', 'www project')
        self.fs.www = www
        self.fs.project = project
        self.argv = []
        self.cookies = SimpleCookie()
        self.short_circuit = True
        self._website = None

    # HTTP Methods
    # ============

    @property
    def website(self):
        if self._website is None:
            argv = [ '--www_root', self.fs.www.root
                   , '--project_root', self.fs.project.root
                   , '--show_tracebacks', '1'
                    ] + self.argv
            self._website = Website(argv)
            self.website.flow.short_circuit = self.short_circuit
        return self._website

    def GET(self, path='/', cookie_info=None, run_through=None, **extra):
        environ = self._build_wsgi_environ(path, "GET", **extra)
        return self._perform_request(environ, cookie_info, run_through)

    def POST(self, path='/', data=None, content_type=MULTIPART_CONTENT, cookie_info=None,
            run_through=None, **extra):
        """Perform a dummy POST request against the test website.

        :param path:
            The url to perform the virutal-POST to.

        :param data:
            A dictionary or list of tuples to be encoded before being POSTed.

        Any additional parameters will be sent as headers. NOTE that in Aspen
        (request.py make_franken_headers) only headers beginning with ``HTTP``
        are included in the request - and those are changed to no longer
        include ``HTTP``. There are currently 2 exceptions to this:
        ``'CONTENT_TYPE'``, ``'CONTENT_LENGTH'`` which are explicitly checked
        for.
        """
        post_data = data if data is not None else {}

        if content_type is MULTIPART_CONTENT:
            post_data = encode_multipart(BOUNDARY, data)

        environ = self._build_wsgi_environ( path
                                          , "POST"
                                          , post_data
                                          , CONTENT_TYPE=str(content_type)
                                          , **extra
                                           )
        return self._perform_request(environ, cookie_info, run_through)

    def simple(self, contents='Greetings, program!', filepath='index.html.spt', urlpath='/',
            run_through=None, want='response', argv=None):
        """A helper to create a file and hit it through our machinery.
        """
        self.fs.www.mk((filepath, contents))
        self.argv = argv if argv is not None else []

        thing = self.GET(urlpath, run_through=run_through)

        attr_path = want.split('.')
        base = attr_path[0]
        attr_path = attr_path[1:]

        if run_through is None:
            out = thing
        else:
            out = thing[base]

        for name in attr_path:
            out = getattr(out, name)

        return out

    # Sockets
    # =======

    def make_request(self, *a, **kw):
        kw['run_through'] = 'dispatch_request_to_filesystem'
        kw['want'] = 'request'
        return self.simple(*a, **kw)

    def make_transport(self, content='', state=0):
        self.fs.www.mk(('echo.sock.spt', content))
        socket = self.make_socket()
        transport = XHRPollingTransport(socket)
        transport.timeout = 0.05 # for testing, could screw up the test
        if state == 1:
            transport.respond(Request(uri='/echo.sock'))
        return transport

    def make_socket_request(self, filename='echo.sock.spt'):
        request = Request(uri='/echo.sock')
        request.website = self.website
        request.fs = self.fs.www.resolve(filename)
        return request

    def make_socket(self, filename='echo.sock.spt', channel=None):
        request = self.make_socket_request(filename='echo.sock.spt')
        if channel is None:
            channel = Channel(request.line.uri.path.raw, ThreadedBuffer)
        socket = Socket(request, channel)
        return socket

    def SocketInThread(harness):

        class _SocketInThread(object):

            def __enter__(self, filename='echo.sock.spt'):
                self.socket = harness.make_socket(filename)
                self.socket.loop.start()
                return self.socket

            def __exit__(self, *a):
                self.socket.loop.stop()

        return _SocketInThread()

    # Hook
    # ====

    def add_cookie_info(self, request, **cookie_info):
        """Place holder function that can be replaced in a subclass.

        For example in gittip.com, it might be of interest to load session
        information into the cookie like this::

            if cookie_info:
                user = cookie_info.get('user')
                if user is not None:
                    user = User.from_id(user)
                    # Note that Cookie needs a bytestring.
                    request.headers.cookie['session'] = user.session_token
        """
        pass

    # Helpers
    # =======

    def _build_wsgi_environ(self, path, method="GET", body=None, **extra):
        environ = {}
        environ['PATH_INFO'] = path
        environ['REMOTE_ADDR'] = b'0.0.0.0'
        environ['REQUEST_METHOD'] = b'GET'
        environ['SERVER_PROTOCOL'] = b'HTTP/1.1'
        environ['HTTP_HOST'] = b'localhost'
        environ['REQUEST_METHOD'] = method
        environ['wsgi.input'] = StringIO(body)
        environ['HTTP_COOKIE'] = self.cookies.output(header='', sep='; ')
        environ.update(extra)
        return environ

    def _perform_request(self, environ, cookie_info, run_through):
        self.add_cookie_info(environ, **(cookie_info or {}))
        out = self.website.respond(environ, _run_through=run_through)
        if run_through is None:
            response = out
        else:
            response = out['response']
        if response is not None:
            if response.headers.cookie:
                self.cookies.update(response.headers.cookie)
        return out
<commit_msg>Remove errant import<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from collections import namedtuple
from Cookie import SimpleCookie
from StringIO import StringIO

from aspen.network_engines import ThreadedBuffer
from aspen.http.request import Request
from aspen.sockets.channel import Channel
from aspen.sockets.socket import Socket
from aspen.sockets.transport import XHRPollingTransport
from aspen.website import Website

BOUNDARY = 'BoUnDaRyStRiNg'
MULTIPART_CONTENT = 'multipart/form-data; boundary=%s' % BOUNDARY

def encode_multipart(boundary, data):
    """
    Encodes multipart POST data from a dictionary of form values.

    Borrowed from Django
    The key will be used as the form data name; the value will be transmitted
    as content. If the value is a file, the contents of the file will be sent
    as an application/octet-stream; otherwise, str(value) will be sent.
    """
    lines = []

    for (key, value) in data.items():
        lines.extend([
            '--' + boundary,
            'Content-Disposition: form-data; name="%s"' % str(key),
            '',
            str(value)
        ])

    lines.extend([
        '--' + boundary + '--',
        '',
    ])
    return '\r\n'.join(lines)

class Harness(object):
    """
    The Aspen testing harness.

    Used in tests to emulate ``GET`` and ``POST`` requests by sending them
    into a ``Website`` instance's ``respond`` method.

    Aspen does not define any User data structures or modules. If you want to
    do anything with users/sessions etc in your tests it is expected that you
    will subclass this class and add a ``add_cookie_info`` method.

    For example, in gittip a suitable subclass might be::

        class GittipTestClient(TestClient):

            def add_cookie_info(self, request, cookie_info):
                if cookie_info:
                    user = cookie_info.get('user')
                    if user is not None:
                        user = User.from_id(user)
                        # Note that Cookie needs a bytestring.
                        request.headers.cookie['session'] = user.session_token

    Example usage in a test::

        def test_api_handles_posts():
            client = TestClient(website)

            # We need to get ourselves a token!
            response = client.get('/')
            csrf_token = response.request.context['csrf_token']

            # Then, add a $1.50 and $3.00 tip
            response = client.post("/test_tippee1/tip.json",
                                {'amount': "1.00", 'csrf_token': csrf_token},
                                cookie_info={'user': 'test_tipper'})

            # Confirm we get back the right amounts in the JSON body.
            first_data = json.loads(response.body)
            assert_equal(first_data['amount'], "1.00")
    """

    def __init__(self, www, project):
        self.fs = namedtuple('fs', 'www project')
        self.fs.www = www
        self.fs.project = project
        self.argv = []
        self.cookies = SimpleCookie()
        self.short_circuit = True
        self._website = None

    # HTTP Methods
    # ============

    @property
    def website(self):
        if self._website is None:
            argv = [ '--www_root', self.fs.www.root
                   , '--project_root', self.fs.project.root
                   , '--show_tracebacks', '1'
                    ] + self.argv
            self._website = Website(argv)
            self.website.flow.short_circuit = self.short_circuit
        return self._website

    def GET(self, path='/', cookie_info=None, run_through=None, **extra):
        environ = self._build_wsgi_environ(path, "GET", **extra)
        return self._perform_request(environ, cookie_info, run_through)

    def POST(self, path='/', data=None, content_type=MULTIPART_CONTENT, cookie_info=None,
            run_through=None, **extra):
        """Perform a dummy POST request against the test website.

        :param path:
            The url to perform the virutal-POST to.

        :param data:
            A dictionary or list of tuples to be encoded before being POSTed.

        Any additional parameters will be sent as headers. NOTE that in Aspen
        (request.py make_franken_headers) only headers beginning with ``HTTP``
        are included in the request - and those are changed to no longer
        include ``HTTP``. There are currently 2 exceptions to this:
        ``'CONTENT_TYPE'``, ``'CONTENT_LENGTH'`` which are explicitly checked
        for.
        """
        post_data = data if data is not None else {}

        if content_type is MULTIPART_CONTENT:
            post_data = encode_multipart(BOUNDARY, data)

        environ = self._build_wsgi_environ( path
                                          , "POST"
                                          , post_data
                                          , CONTENT_TYPE=str(content_type)
                                          , **extra
                                           )
        return self._perform_request(environ, cookie_info, run_through)

    def simple(self, contents='Greetings, program!', filepath='index.html.spt', urlpath='/',
            run_through=None, want='response', argv=None):
        """A helper to create a file and hit it through our machinery.
        """
        self.fs.www.mk((filepath, contents))
        self.argv = argv if argv is not None else []

        thing = self.GET(urlpath, run_through=run_through)

        attr_path = want.split('.')
        base = attr_path[0]
        attr_path = attr_path[1:]

        if run_through is None:
            out = thing
        else:
            out = thing[base]

        for name in attr_path:
            out = getattr(out, name)

        return out

    # Sockets
    # =======

    def make_request(self, *a, **kw):
        kw['run_through'] = 'dispatch_request_to_filesystem'
        kw['want'] = 'request'
        return self.simple(*a, **kw)

    def make_transport(self, content='', state=0):
        self.fs.www.mk(('echo.sock.spt', content))
        socket = self.make_socket()
        transport = XHRPollingTransport(socket)
        transport.timeout = 0.05 # for testing, could screw up the test
        if state == 1:
            transport.respond(Request(uri='/echo.sock'))
        return transport

    def make_socket_request(self, filename='echo.sock.spt'):
        request = Request(uri='/echo.sock')
        request.website = self.website
        request.fs = self.fs.www.resolve(filename)
        return request

    def make_socket(self, filename='echo.sock.spt', channel=None):
        request = self.make_socket_request(filename='echo.sock.spt')
        if channel is None:
            channel = Channel(request.line.uri.path.raw, ThreadedBuffer)
        socket = Socket(request, channel)
        return socket

    def SocketInThread(harness):

        class _SocketInThread(object):

            def __enter__(self, filename='echo.sock.spt'):
                self.socket = harness.make_socket(filename)
                self.socket.loop.start()
                return self.socket

            def __exit__(self, *a):
                self.socket.loop.stop()

        return _SocketInThread()

    # Hook
    # ====

    def add_cookie_info(self, request, **cookie_info):
        """Place holder function that can be replaced in a subclass.

        For example in gittip.com, it might be of interest to load session
        information into the cookie like this::

            if cookie_info:
                user = cookie_info.get('user')
                if user is not None:
                    user = User.from_id(user)
                    # Note that Cookie needs a bytestring.
                    request.headers.cookie['session'] = user.session_token
        """
        pass

    # Helpers
    # =======

    def _build_wsgi_environ(self, path, method="GET", body=None, **extra):
        environ = {}
        environ['PATH_INFO'] = path
        environ['REMOTE_ADDR'] = b'0.0.0.0'
        environ['REQUEST_METHOD'] = b'GET'
        environ['SERVER_PROTOCOL'] = b'HTTP/1.1'
        environ['HTTP_HOST'] = b'localhost'
        environ['REQUEST_METHOD'] = method
        environ['wsgi.input'] = StringIO(body)
        environ['HTTP_COOKIE'] = self.cookies.output(header='', sep='; ')
        environ.update(extra)
        return environ

    def _perform_request(self, environ, cookie_info, run_through):
        self.add_cookie_info(environ, **(cookie_info or {}))
        out = self.website.respond(environ, _run_through=run_through)
        if run_through is None:
            response = out
        else:
            response = out['response']
        if response is not None:
            if response.headers.cookie:
                self.cookies.update(response.headers.cookie)
        return out
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""\
This is a python port of "Goose" orignialy licensed to Gravity.com
under one or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.

Python port was written by Xavier Grangier for Recrutae

Gravity.com licenses this file
to you 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.
"""
import os
import glob
from copy import deepcopy
from goose.article import Article
from goose.utils import URLHelper, RawHelper
from goose.extractors import StandardContentExtractor
from goose.cleaners import StandardDocumentCleaner
from goose.outputformatters import StandardOutputFormatter
from goose.images.extractors import UpgradedImageIExtractor
from goose.videos.extractors import VideoExtractor
from goose.network import HtmlFetcher

class CrawlCandidate(object):

    def __init__(self, config, url, raw_html):
        self.config = config
        # parser
        self.parser = self.config.get_parser()
        self.url = url
        self.raw_html = raw_html

class Crawler(object):

    def __init__(self, config):
        # config
        self.config = config
        # parser
        self.parser = self.config.get_parser()

        # article
        self.article = Article()

        # init the extractor
        self.extractor = self.get_extractor()

        # init the document cleaner
        self.cleaner = self.get_cleaner()

        # init the output formatter
        self.formatter = self.get_formatter()

        # video extractor
        self.video_extractor = self.get_video_extractor()

        # image extrator
        self.image_extractor = self.get_image_extractor()

        # html fetcher
        self.htmlfetcher = HtmlFetcher(self.config)

        # TODO : log prefix
        self.logPrefix = "crawler:"

    def crawl(self, crawl_candidate):

        # parser candidate
        parse_candidate = self.get_parse_candidate(crawl_candidate)

        # raw html
        raw_html = self.get_html(crawl_candidate, parse_candidate)

        if raw_html is None:
            return self.article

        # create document
        doc = self.get_document(raw_html)

        # article
        self.article.final_url = parse_candidate.url
        self.article.link_hash = parse_candidate.link_hash
        self.article.raw_html = raw_html
        self.article.doc = doc
        self.article.raw_doc = deepcopy(doc)
        self.article.publish_date = self.extractor.get_publish_date()
        # self.article.additional_data = config.get_additionaldata_extractor.extract(doc)
        self.article.title = self.extractor.get_title()
        self.article.meta_lang = self.extractor.get_meta_lang()
        self.article.meta_favicon = self.extractor.get_favicon()
        self.article.meta_description = self.extractor.get_meta_description()
        self.article.meta_keywords = self.extractor.get_meta_keywords()
        self.article.canonical_link = self.extractor.get_canonical_link()
        self.article.domain = self.extractor.get_domain()
        self.article.tags = self.extractor.extract_tags()

        # before we do any calcs on the body itself let's clean up the document
        self.article.doc = self.cleaner.clean()

        # big stuff
        self.article.top_node = self.extractor.calculate_best_node()

        # if we have a top node
        # let's process it
        if self.article.top_node is not None:

            # video handeling
            self.video_extractor.get_videos()

            # image handeling
            if self.config.enable_image_fetching:
                self.get_image()

            # post cleanup
            self.article.top_node = self.extractor.post_cleanup()

            # clean_text
            self.article.cleaned_text = self.formatter.get_formatted_text()

        # cleanup tmp file
        self.relase_resources()

        # return the article
        return self.article

    def get_parse_candidate(self, crawl_candidate):
        if crawl_candidate.raw_html:
            return RawHelper.get_parsing_candidate(crawl_candidate.url, crawl_candidate.raw_html)
        return URLHelper.get_parsing_candidate(crawl_candidate.url)

    def get_image(self):
        doc = self.article.raw_doc
        top_node = self.article.top_node
        self.article.top_image = self.image_extractor.get_best_image(doc, top_node)

    def get_html(self, crawl_candidate, parsing_candidate):
        # we got a raw_tml
        # no need to fetch remote content
        if crawl_candidate.raw_html:
            return crawl_candidate.raw_html

        # fetch HTML
        html = self.htmlfetcher.get_html(parsing_candidate.url)
        self.article.additional_data.update({
            'request': self.htmlfetcher.request,
            'result': self.htmlfetcher.result,
            })
        return html

    def get_image_extractor(self):
        return UpgradedImageIExtractor(self.config, self.article)

    def get_video_extractor(self):
        return VideoExtractor(self.config, self.article)

    def get_formatter(self):
        return StandardOutputFormatter(self.config, self.article)

    def get_cleaner(self):
        return StandardDocumentCleaner(self.config, self.article)

    def get_document(self, raw_html):
        doc = self.parser.fromstring(raw_html)
        return doc

    def get_extractor(self):
        return StandardContentExtractor(self.config, self.article)

    def relase_resources(self):
        path = os.path.join(self.config.local_storage_path, '%s_*' % self.article.link_hash)
        for fname in glob.glob(path):
            try:
                os.remove(fname)
            except OSError:
                # TODO better log handeling
                pass
<commit_msg>#165 - extract opengraph data<commit_after># -*- coding: utf-8 -*-
"""\
This is a python port of "Goose" orignialy licensed to Gravity.com
under one or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.

Python port was written by Xavier Grangier for Recrutae

Gravity.com licenses this file
to you 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.
"""
import os
import glob
from copy import deepcopy
from goose.article import Article
from goose.utils import URLHelper, RawHelper
from goose.extractors import StandardContentExtractor
from goose.cleaners import StandardDocumentCleaner
from goose.outputformatters import StandardOutputFormatter
from goose.images.extractors import UpgradedImageIExtractor
from goose.videos.extractors import VideoExtractor
from goose.network import HtmlFetcher

class CrawlCandidate(object):

    def __init__(self, config, url, raw_html):
        self.config = config
        # parser
        self.parser = self.config.get_parser()
        self.url = url
        self.raw_html = raw_html

class Crawler(object):

    def __init__(self, config):
        # config
        self.config = config
        # parser
        self.parser = self.config.get_parser()

        # article
        self.article = Article()

        # init the extractor
        self.extractor = self.get_extractor()

        # init the document cleaner
        self.cleaner = self.get_cleaner()

        # init the output formatter
        self.formatter = self.get_formatter()

        # video extractor
        self.video_extractor = self.get_video_extractor()

        # image extrator
        self.image_extractor = self.get_image_extractor()

        # html fetcher
        self.htmlfetcher = HtmlFetcher(self.config)

        # TODO : log prefix
        self.logPrefix = "crawler:"

    def crawl(self, crawl_candidate):

        # parser candidate
        parse_candidate = self.get_parse_candidate(crawl_candidate)

        # raw html
        raw_html = self.get_html(crawl_candidate, parse_candidate)

        if raw_html is None:
            return self.article

        # create document
        doc = self.get_document(raw_html)

        # article
        self.article.final_url = parse_candidate.url
        self.article.link_hash = parse_candidate.link_hash
        self.article.raw_html = raw_html
        self.article.doc = doc
        self.article.raw_doc = deepcopy(doc)
        self.article.publish_date = self.extractor.get_publish_date()
        # self.article.additional_data = config.get_additionaldata_extractor.extract(doc)
        self.article.title = self.extractor.get_title()
        self.article.meta_lang = self.extractor.get_meta_lang()
        self.article.meta_favicon = self.extractor.get_favicon()
        self.article.meta_description = self.extractor.get_meta_description()
        self.article.meta_keywords = self.extractor.get_meta_keywords()
        self.article.canonical_link = self.extractor.get_canonical_link()
        self.article.domain = self.extractor.get_domain()
        self.article.tags = self.extractor.extract_tags()

        # opengraph
        self.article.opengraph = self.extractor.extract_opengraph()

        # before we do any calcs on the body itself let's clean up the document
        self.article.doc = self.cleaner.clean()

        # big stuff
        self.article.top_node = self.extractor.calculate_best_node()

        # if we have a top node
        # let's process it
        if self.article.top_node is not None:

            # video handeling
            self.video_extractor.get_videos()

            # image handeling
            if self.config.enable_image_fetching:
                self.get_image()

            # post cleanup
            self.article.top_node = self.extractor.post_cleanup()

            # clean_text
            self.article.cleaned_text = self.formatter.get_formatted_text()

        # cleanup tmp file
        self.relase_resources()

        # return the article
        return self.article

    def get_parse_candidate(self, crawl_candidate):
        if crawl_candidate.raw_html:
            return RawHelper.get_parsing_candidate(crawl_candidate.url, crawl_candidate.raw_html)
        return URLHelper.get_parsing_candidate(crawl_candidate.url)

    def get_image(self):
        doc = self.article.raw_doc
        top_node = self.article.top_node
        self.article.top_image = self.image_extractor.get_best_image(doc, top_node)

    def get_html(self, crawl_candidate, parsing_candidate):
        # we got a raw_tml
        # no need to fetch remote content
        if crawl_candidate.raw_html:
            return crawl_candidate.raw_html

        # fetch HTML
        html = self.htmlfetcher.get_html(parsing_candidate.url)
        self.article.additional_data.update({
            'request': self.htmlfetcher.request,
            'result': self.htmlfetcher.result,
            })
        return html

    def get_image_extractor(self):
        return UpgradedImageIExtractor(self.config, self.article)

    def get_video_extractor(self):
        return VideoExtractor(self.config, self.article)

    def get_formatter(self):
        return StandardOutputFormatter(self.config, self.article)

    def get_cleaner(self):
        return StandardDocumentCleaner(self.config, self.article)

    def get_document(self, raw_html):
        doc = self.parser.fromstring(raw_html)
        return doc

    def get_extractor(self):
        return StandardContentExtractor(self.config, self.article)

    def relase_resources(self):
        path = os.path.join(self.config.local_storage_path, '%s_*' % self.article.link_hash)
        for fname in glob.glob(path):
            try:
                os.remove(fname)
            except OSError:
                # TODO better log handeling
                pass
<|endoftext|>


<commit_before>"""
Test numpy.timedelta64 support.
"""

from __future__ import print_function

import itertools

import numpy as np

import numba.unittest_support as unittest
from numba.typeinfer import TypingError
from numba import config, jit, npdatetime, types
from .support import TestCase

date_units = ('Y', 'M')
time_units = ('W', 'D', 'h', 'm', 's', 'ms', 'us', 'ns', 'ps', 'fs', 'as')
# All except generic ("")
all_units = date_units + time_units

def add_usecase(x, y):
    return x + y

def sub_usecase(x, y):
    return x - y

def mul_usecase(x, y):
    return x * y

def div_usecase(x, y):
    return x / y

def floordiv_usecase(x, y):
    return x // y

class TestModuleHelpers(TestCase):
    """
    Test the various helpers in numba.npdatetime.
    """

    def test_can_cast_timedelta(self):
        f = npdatetime.can_cast_timedelta_units
        for a, b in itertools.product(date_units, time_units):
            self.assertFalse(f(a, b), (a, b))
            self.assertFalse(f(b, a), (a, b))
        for unit in all_units:
            self.assertFalse(f(unit, ''))
            self.assertTrue(f('', unit))
        for unit in all_units + ('',):
            self.assertTrue(f(unit, unit))

        def check_units_group(group):
            for i, a in enumerate(group):
                for b in group[:i]:
                    # large into smaller is ok
                    self.assertTrue(f(b, a))
                    # small into larger is not
                    self.assertFalse(f(a, b))

        check_units_group(date_units)
        check_units_group(time_units)

    def test_timedelta_conversion(self):
        f = npdatetime.get_timedelta_conversion_factor
        for unit in all_units + ('',):
            self.assertEqual(f(unit, unit), 1)
        for unit in all_units:
            self.assertEqual(f('', unit), 1)
        for a, b in itertools.product(time_units, date_units):
            self.assertIs(f(a, b), None)
            self.assertIs(f(b, a), None)

        def check_units_group(group):
            for i, a in enumerate(group):
                for b in group[:i]:
                    self.assertGreater(f(b, a), 1, (b, a))
                    self.assertIs(f(a, b), None)

        check_units_group(date_units)
        check_units_group(time_units)

        # Check some hand-picked values
        self.assertEqual(f('Y', 'M'), 12)
        self.assertEqual(f('W', 'h'), 24 * 7)
        self.assertEqual(f('W', 'm'), 24 * 7 * 60)
        self.assertEqual(f('W', 'us'), 24 * 7 * 3600 * 1000 * 1000)

TD = np.timedelta64
DT = np.datetime64

class TestScalarOperators(TestCase):

    jitargs = dict()

    def jit(self, pyfunc):
        return jit(**self.jitargs)(pyfunc)

    def test_add(self):
        f = self.jit(add_usecase)
        def check(a, b, expected):
            self.assertPreciseEqual(f(a, b), expected)
            self.assertPreciseEqual(f(b, a), expected)

        check(TD(1), TD(2), TD(3))
        check(TD(1, 's'), TD(2, 's'), TD(3, 's'))
        # Implicit unit promotion
        check(TD(1), TD(2, 's'), TD(3, 's'))
        check(TD(1), TD(2, 'ms'), TD(3, 'ms'))
        check(TD(1, 's'), TD(2, 'us'), TD(1000002, 'us'))
        check(TD(1, 'W'), TD(2, 'D'), TD(9, 'D'))
        # NaTs
        check(TD('NaT'), TD(1), TD('NaT'))
        check(TD('NaT', 's'), TD(1, 'D'), TD('NaT', 's'))
        check(TD('NaT', 's'), TD(1, 'ms'), TD('NaT', 'ms'))
        # Cannot add days and weeks
        with self.assertRaises((TypeError, TypingError)):
            f(TD(1, 'M'), TD(1, 'D'))

    def test_sub(self):
        f = self.jit(sub_usecase)
        def check(a, b, expected):
            self.assertPreciseEqual(f(a, b), expected)
            self.assertPreciseEqual(f(b, a), -expected)

        check(TD(3), TD(2), TD(1))
        check(TD(3, 's'), TD(2, 's'), TD(1, 's'))
        # Implicit unit promotion
        check(TD(3), TD(2, 's'), TD(1, 's'))
        check(TD(3), TD(2, 'ms'), TD(1, 'ms'))
        check(TD(3, 's'), TD(2, 'us'), TD(2999998, 'us'))
        check(TD(1, 'W'), TD(2, 'D'), TD(5, 'D'))
        # NaTs
        check(TD('NaT'), TD(1), TD('NaT'))
        check(TD('NaT', 's'), TD(1, 'D'), TD('NaT', 's'))
        check(TD('NaT', 's'), TD(1, 'ms'), TD('NaT', 'ms'))
        # Cannot sub days to weeks
        with self.assertRaises((TypeError, TypingError)):
            f(TD(1, 'M'), TD(1, 'D'))

    def test_mul(self):
        f = self.jit(mul_usecase)
        def check(a, b, expected):
            self.assertPreciseEqual(f(a, b), expected)
            self.assertPreciseEqual(f(b, a), expected)

        # int * timedelta64
        check(TD(3), 2, TD(6))
        check(TD(3, 'ps'), 2, TD(6, 'ps'))
        check(TD('NaT', 'ps'), 2, TD('NaT', 'ps'))
        # float * timedelta64
        check(TD(7), 1.5, TD(10))
        check(TD(-7), 1.5, TD(-10))
        check(TD(7, 'ps'), -1.5, TD(-10, 'ps'))
        check(TD(-7), -1.5, TD(10))
        check(TD('NaT', 'ps'), -1.5, TD('NaT', 'ps'))
        check(TD(7, 'ps'), float('nan'), TD('NaT', 'ps'))
        # wraparound on overflow
        check(TD(2**62, 'ps'), 16, TD(0, 'ps'))

    def test_div(self):
        div = self.jit(div_usecase)
        floordiv = self.jit(floordiv_usecase)
        def check(a, b, expected):
            self.assertPreciseEqual(div(a, b), expected)
            self.assertPreciseEqual(floordiv(a, b), expected)

        # timedelta64 / int
        check(TD(3), 2, TD(1))
        check(TD(-3, 'ps'), 2, TD(-1, 'ps'))
        check(TD('NaT', 'ps'), 2, TD('NaT', 'ps'))
        check(TD(3, 'ps'), 0, TD('NaT', 'ps'))
        check(TD('NaT', 'ps'), 0, TD('NaT', 'ps'))
        # timedelta64 / float
        check(TD(7), 0.5, TD(14))
        check(TD(-7, 'ps'), 1.5, TD(-4, 'ps'))
        check(TD('NaT', 'ps'), 2.5, TD('NaT', 'ps'))
        check(TD(3, 'ps'), 0.0, TD('NaT', 'ps'))
        check(TD('NaT', 'ps'), 0.0, TD('NaT', 'ps'))
        check(TD(3, 'ps'), float('nan'), TD('NaT', 'ps'))
        check(TD('NaT', 'ps'), float('nan'), TD('NaT', 'ps'))

    def test_homogenous_div(self):
        div = self.jit(div_usecase)
        def check(a, b, expected):
            self.assertPreciseEqual(div(a, b), expected)

        # timedelta64 / timedelta64
        check(TD(7), TD(3), 7. / 3.)
        check(TD(7), TD(3, 'ms'), 7. / 3.)
        check(TD(7, 'us'), TD(3, 'ms'), 7. / 3000.)
        check(TD(7, 'ms'), TD(3, 'us'), 7000. / 3.)
        check(TD(7), TD(0), float('+inf'))
        check(TD(-7), TD(0), float('-inf'))
        # NaTs
        check(TD('nat'), TD(3), float('nan'))
        check(TD(3), TD('nat'), float('nan'))
        check(TD('nat'), TD(0), float('nan'))

class TestScalarOperatorsNoPython(TestScalarOperators):

    jitargs = dict(nopython=True)

if __name__ == '__main__':
    unittest.main()
<commit_msg>Add a test for 0 / 0<commit_after>"""
Test numpy.timedelta64 support.
"""

from __future__ import print_function

import itertools

import numpy as np

import numba.unittest_support as unittest
from numba.typeinfer import TypingError
from numba import config, jit, npdatetime, types
from .support import TestCase

date_units = ('Y', 'M')
time_units = ('W', 'D', 'h', 'm', 's', 'ms', 'us', 'ns', 'ps', 'fs', 'as')
# All except generic ("")
all_units = date_units + time_units

def add_usecase(x, y):
    return x + y

def sub_usecase(x, y):
    return x - y

def mul_usecase(x, y):
    return x * y

def div_usecase(x, y):
    return x / y

def floordiv_usecase(x, y):
    return x // y

class TestModuleHelpers(TestCase):
    """
    Test the various helpers in numba.npdatetime.
    """

    def test_can_cast_timedelta(self):
        f = npdatetime.can_cast_timedelta_units
        for a, b in itertools.product(date_units, time_units):
            self.assertFalse(f(a, b), (a, b))
            self.assertFalse(f(b, a), (a, b))
        for unit in all_units:
            self.assertFalse(f(unit, ''))
            self.assertTrue(f('', unit))
        for unit in all_units + ('',):
            self.assertTrue(f(unit, unit))

        def check_units_group(group):
            for i, a in enumerate(group):
                for b in group[:i]:
                    # large into smaller is ok
                    self.assertTrue(f(b, a))
                    # small into larger is not
                    self.assertFalse(f(a, b))

        check_units_group(date_units)
        check_units_group(time_units)

    def test_timedelta_conversion(self):
        f = npdatetime.get_timedelta_conversion_factor
        for unit in all_units + ('',):
            self.assertEqual(f(unit, unit), 1)
        for unit in all_units:
            self.assertEqual(f('', unit), 1)
        for a, b in itertools.product(time_units, date_units):
            self.assertIs(f(a, b), None)
            self.assertIs(f(b, a), None)

        def check_units_group(group):
            for i, a in enumerate(group):
                for b in group[:i]:
                    self.assertGreater(f(b, a), 1, (b, a))
                    self.assertIs(f(a, b), None)

        check_units_group(date_units)
        check_units_group(time_units)

        # Check some hand-picked values
        self.assertEqual(f('Y', 'M'), 12)
        self.assertEqual(f('W', 'h'), 24 * 7)
        self.assertEqual(f('W', 'm'), 24 * 7 * 60)
        self.assertEqual(f('W', 'us'), 24 * 7 * 3600 * 1000 * 1000)

TD = np.timedelta64
DT = np.datetime64

class TestScalarOperators(TestCase):

    jitargs = dict()

    def jit(self, pyfunc):
        return jit(**self.jitargs)(pyfunc)

    def test_add(self):
        f = self.jit(add_usecase)
        def check(a, b, expected):
            self.assertPreciseEqual(f(a, b), expected)
            self.assertPreciseEqual(f(b, a), expected)

        check(TD(1), TD(2), TD(3))
        check(TD(1, 's'), TD(2, 's'), TD(3, 's'))
        # Implicit unit promotion
        check(TD(1), TD(2, 's'), TD(3, 's'))
        check(TD(1), TD(2, 'ms'), TD(3, 'ms'))
        check(TD(1, 's'), TD(2, 'us'), TD(1000002, 'us'))
        check(TD(1, 'W'), TD(2, 'D'), TD(9, 'D'))
        # NaTs
        check(TD('NaT'), TD(1), TD('NaT'))
        check(TD('NaT', 's'), TD(1, 'D'), TD('NaT', 's'))
        check(TD('NaT', 's'), TD(1, 'ms'), TD('NaT', 'ms'))
        # Cannot add days and weeks
        with self.assertRaises((TypeError, TypingError)):
            f(TD(1, 'M'), TD(1, 'D'))

    def test_sub(self):
        f = self.jit(sub_usecase)
        def check(a, b, expected):
            self.assertPreciseEqual(f(a, b), expected)
            self.assertPreciseEqual(f(b, a), -expected)

        check(TD(3), TD(2), TD(1))
        check(TD(3, 's'), TD(2, 's'), TD(1, 's'))
        # Implicit unit promotion
        check(TD(3), TD(2, 's'), TD(1, 's'))
        check(TD(3), TD(2, 'ms'), TD(1, 'ms'))
        check(TD(3, 's'), TD(2, 'us'), TD(2999998, 'us'))
        check(TD(1, 'W'), TD(2, 'D'), TD(5, 'D'))
        # NaTs
        check(TD('NaT'), TD(1), TD('NaT'))
        check(TD('NaT', 's'), TD(1, 'D'), TD('NaT', 's'))
        check(TD('NaT', 's'), TD(1, 'ms'), TD('NaT', 'ms'))
        # Cannot sub days to weeks
        with self.assertRaises((TypeError, TypingError)):
            f(TD(1, 'M'), TD(1, 'D'))

    def test_mul(self):
        f = self.jit(mul_usecase)
        def check(a, b, expected):
            self.assertPreciseEqual(f(a, b), expected)
            self.assertPreciseEqual(f(b, a), expected)

        # int * timedelta64
        check(TD(3), 2, TD(6))
        check(TD(3, 'ps'), 2, TD(6, 'ps'))
        check(TD('NaT', 'ps'), 2, TD('NaT', 'ps'))
        # float * timedelta64
        check(TD(7), 1.5, TD(10))
        check(TD(-7), 1.5, TD(-10))
        check(TD(7, 'ps'), -1.5, TD(-10, 'ps'))
        check(TD(-7), -1.5, TD(10))
        check(TD('NaT', 'ps'), -1.5, TD('NaT', 'ps'))
        check(TD(7, 'ps'), float('nan'), TD('NaT', 'ps'))
        # wraparound on overflow
        check(TD(2**62, 'ps'), 16, TD(0, 'ps'))

    def test_div(self):
        div = self.jit(div_usecase)
        floordiv = self.jit(floordiv_usecase)
        def check(a, b, expected):
            self.assertPreciseEqual(div(a, b), expected)
            self.assertPreciseEqual(floordiv(a, b), expected)

        # timedelta64 / int
        check(TD(3), 2, TD(1))
        check(TD(-3, 'ps'), 2, TD(-1, 'ps'))
        check(TD('NaT', 'ps'), 2, TD('NaT', 'ps'))
        check(TD(3, 'ps'), 0, TD('NaT', 'ps'))
        check(TD('NaT', 'ps'), 0, TD('NaT', 'ps'))
        # timedelta64 / float
        check(TD(7), 0.5, TD(14))
        check(TD(-7, 'ps'), 1.5, TD(-4, 'ps'))
        check(TD('NaT', 'ps'), 2.5, TD('NaT', 'ps'))
        check(TD(3, 'ps'), 0.0, TD('NaT', 'ps'))
        check(TD('NaT', 'ps'), 0.0, TD('NaT', 'ps'))
        check(TD(3, 'ps'), float('nan'), TD('NaT', 'ps'))
        check(TD('NaT', 'ps'), float('nan'), TD('NaT', 'ps'))

    def test_homogenous_div(self):
        div = self.jit(div_usecase)
        def check(a, b, expected):
            self.assertPreciseEqual(div(a, b), expected)

        # timedelta64 / timedelta64
        check(TD(7), TD(3), 7. / 3.)
        check(TD(7), TD(3, 'ms'), 7. / 3.)
        check(TD(7, 'us'), TD(3, 'ms'), 7. / 3000.)
        check(TD(7, 'ms'), TD(3, 'us'), 7000. / 3.)
        check(TD(7), TD(0), float('+inf'))
        check(TD(-7), TD(0), float('-inf'))
        check(TD(0), TD(0), float('nan'))
        # NaTs
        check(TD('nat'), TD(3), float('nan'))
        check(TD(3), TD('nat'), float('nan'))
        check(TD('nat'), TD(0), float('nan'))

class TestScalarOperatorsNoPython(TestScalarOperators):

    jitargs = dict(nopython=True)

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


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright (C) 2017, AGB & GC
# Full license can be found in License.md
#-----------------------------------------------------------------------------
""" Functions that specify the boundary location as a function of MLT

Functions
-------------------------------------------------------------------------------
circular(ocb, mlt, r_add)
    Return a circular boundary
ampere_harmonic(ocb, mlt)
    Return the results of a 3rd order harmonic fit to AMPERE/DMSP differences

References
-------------------------------------------------------------------------------
Burrell paper in prep

"""
import numpy as np

def circular(ocb, mlt, mlt_coords="aacgm", r_add=0.0):
    """Return a circular boundary

    Parameters
    ----------
    ocb : (OCBoundary)
        OCBoundary object
    mlt : (float)
        Magnetic local time in hours (not actually used)
    mlt_coords : (str)
        Coordinates of the MLT.  Accepts 'aacgm' and 'ocb' (default='aacgm')
        (not used here)
    r_add : (float)
        Offset added to default radius (default=0.0)

    Returns
    -------
    r : (float)
        Radius at this MLT and time (as specified by ocb.rec_ind)

    """

    if ocb.rec_ind < 0 or ocb.rec_ind >= ocb.records:
        r = np.nan
    else:
        r = ocb.r[ocb.rec_ind] + r_add

    return r

def ampere_harmonic(ocb, mlt, mlt_coords="aacgm", method='median'):
    """Return the results of a 3rd order harmonic fit to AMPERE/DMSP differences

    Parameters
    ----------
    ocb : (OCBoundary)
        OCBoundary object
    mlt : (float or array-like)
        Magnetic local time in hours (not actually used)
    mlt_coords : (str)
        Coordinates of the MLT.  Accepts 'aacgm' and 'ocb' (default='aacgm')
    method : (str)
        Method used to determine coefficients; accepts median or gaussian
        (default='median')

    Returns
    -------
    r : (float)
        Radius at this MLT and time (as specified by ocb.rec_ind)

    """
    from ocbpy.ocb_time import hr2rad

    mlt_coods = mlt_coords.lower()

    if mlt_coords == "ocb":
        raise ValueError("routine lacks OCB to AACGM conversion")

    coeff = {'median': [3.10, -0.10, -0.08, 1.83, -0.73, -0.80, -0.07, 0.28,
                        -0.67, 0.93, -0.12, -0.21, 0.08],
             'gaussian': [3.61, -0.19, 0.15, 1.74, -0.67, -0.85, -0.13, 0.26,
                          -0.42, 0.86, 0.05, 0.27, 0.12]}

    if method not in coeff.keys():
        raise ValueError("unknown coefficient computation method")

    if ocb.rec_ind < 0 or ocb.rec_ind >= ocb.records:
        r = np.nan
    else:
        rad_mlt = hr2rad(mlt)
        r_add = coeff[method][0] \
            + coeff[method][1] * np.cos(rad_mlt+coeff[method][2]) \
            + coeff[method][3] * np.sin(rad_mlt+coeff[method][4]) \
            + coeff[method][5] * np.cos(2.0 * (rad_mlt+coeff[method][6])) \
            + coeff[method][7] * np.sin(2.0 * (rad_mlt+coeff[method][8])) \
            + coeff[method][9] * np.cos(3.0 * (rad_mlt+coeff[method][10])) \
            + coeff[method][11] * np.sin(3.0 * (rad_mlt+coeff[method][12]))
        r = ocb.r[ocb.rec_ind] + r_add

    return r
<commit_msg>updated correction function method<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright (C) 2017, AGB & GC
# Full license can be found in License.md
#-----------------------------------------------------------------------------
""" Functions that specify the boundary location as a function of MLT

Functions
-------------------------------------------------------------------------------
circular(ocb, mlt, r_add)
    Return a circular boundary
ampere_harmonic(ocb, mlt)
    Return the results of a 3rd order harmonic fit to AMPERE/DMSP differences

References
-------------------------------------------------------------------------------
Burrell paper in prep

"""
import numpy as np

def circular(mlt, r_add=0.0):
    """Return a circular boundary correction

    Parameters
    ----------
    mlt : (float)
        Magnetic local time in hours (not actually used)
    r_add : (float)
        Offset added to default radius in degrees (default=0.0)

    Returns
    -------
    r_corr : (float)
        Radius correction in degrees at this MLT

    """

    return r_add

def ampere_harmonic(mlt, method='median'):
    """Return the results of a 3rd order harmonic fit to AMPERE/DMSP differences

    Parameters
    ----------
    mlt : (float or array-like)
        Magnetic local time in hours
    method : (str)
        Method used to determine coefficients; accepts median or gaussian
        (default='median')

    Returns
    -------
    r_corr : (float)
        Radius correction in degrees at this MLT

    """
    from ocbpy.ocb_time import hr2rad

    method = method.lower()
    coeff = {'median': [3.10, -0.10, -0.08, 1.83, -0.73, -0.80, -0.07, 0.28,
                        -0.67, 0.93, -0.12, -0.21, 0.08],
             'gaussian': [3.61, -0.19, 0.15, 1.74, -0.67, -0.85, -0.13, 0.26,
                          -0.42, 0.86, 0.05, 0.27, 0.12]}

    if method not in coeff.keys():
        raise ValueError("unknown coefficient computation method")

    rad_mlt = hr2rad(mlt)
    r_corr = coeff[method][0] \
                + coeff[method][1] * np.cos(rad_mlt+coeff[method][2]) \
                + coeff[method][3] * np.sin(rad_mlt+coeff[method][4]) \
                + coeff[method][5] * np.cos(2.0 * (rad_mlt+coeff[method][6])) \
                + coeff[method][7] * np.sin(2.0 * (rad_mlt+coeff[method][8])) \
                + coeff[method][9] * np.cos(3.0 * (rad_mlt+coeff[method][10])) \
                + coeff[method][11] * np.sin(3.0 * (rad_mlt+coeff[method][12]))

    return r_corr
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# Copyright (C) 2011 Romain Bignon, Laurent Bachelier
#
# This file is part of ass2m.
#
# ass2m is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# ass2m 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with ass2m. If not, see <http://www.gnu.org/licenses/>.

import os
from PIL import Image
from paste.httpheaders import CACHE_CONTROL

from ass2m.plugin import Plugin
from ass2m.routes import View
from ass2m.server import ViewAction, FileApp

class ListGalleryAction(ViewAction):
    def get(self):
        dirs = []
        photos = []
        description = None
        for f in self.ctx.iter_files():
            if f.isdir():
                dirs.append(f)
            elif self.ctx.user.has_perms(f, f.PERM_READ):
                filename = f.get_name()
                mimetype = f.get_mimetype()
                if mimetype is not None and mimetype.startswith('image'):
                    photos.append(f)
                elif filename == 'DESCRIPTION':
                    with open(f.get_realpath(), 'r') as f:
                        description = f.read()

        self.ctx.template_vars['dirs'] = dirs
        self.ctx.template_vars['photos'] = photos
        self.ctx.template_vars['description'] = description
        self.ctx.res.body = self.ctx.render('list-gallery.html')

class DownloadThumbnailAction(ViewAction):
    DEFAULT_SIZE = 300

    def _get_size(self):
        try:
            size = int(self.ctx.req.str_GET['thumb_size'])
        except (KeyError,ValueError):
            size = self.DEFAULT_SIZE
        else:
            if size < 1 or size > 1000:
                size = self.DEFAULT_SIZE
        return size

    def get(self):
        size = self._get_size()
        thumbdir = os.path.join(self.ctx.storage.path, 'thumbnails')
        f = self.ctx.file

        # use a lossless format in doubt
        lossy = f.get_mimetype() == 'image/jpeg'
        thumbext = 'jpg' if lossy else 'png'
        thumbpath = os.path.join(thumbdir, str(size), '%s.%s' % (f.get_hash(), thumbext))

        mtime = int(os.path.getmtime(f.get_realpath()))
        if not os.path.exists(thumbpath) or mtime > int(os.path.getmtime(thumbpath)):
            if not os.path.isdir(os.path.dirname(thumbpath)):
                os.makedirs(os.path.dirname(thumbpath))
            with open(f.get_realpath(), 'rb') as fp:
                img = Image.open(fp)
                img.thumbnail((size, size), Image.BILINEAR)
                with open(thumbpath, 'wb') as tfp:
                    if lossy:
                        img.save(tfp, 'jpeg', quality=95)
                    else:
                        img.save(tfp, 'png')
            os.utime(thumbpath, (mtime, mtime))

        self.ctx.res = FileApp(thumbpath)
        self.ctx.res.cache_control(private=True, max_age=CACHE_CONTROL.ONE_HOUR)

class GalleryPlugin(Plugin):
    def init(self):
        self.register_web_view(
            View(object_type='directory', name='gallery'),
            ListGalleryAction)
        self.register_web_view(
            View(object_type='file', mimetype='image', name='thumbnail'),
            DownloadThumbnailAction)
<commit_msg>Store metadata to ease thumbnail garbage collection<commit_after># -*- coding: utf-8 -*-

# Copyright (C) 2011 Romain Bignon, Laurent Bachelier
#
# This file is part of ass2m.
#
# ass2m is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# ass2m 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with ass2m. If not, see <http://www.gnu.org/licenses/>.

import os
from PIL import Image
from paste.httpheaders import CACHE_CONTROL

from ass2m.plugin import Plugin
from ass2m.routes import View
from ass2m.server import ViewAction, FileApp

class ListGalleryAction(ViewAction):
    def get(self):
        dirs = []
        photos = []
        description = None
        for f in self.ctx.iter_files():
            if f.isdir():
                dirs.append(f)
            elif self.ctx.user.has_perms(f, f.PERM_READ):
                filename = f.get_name()
                mimetype = f.get_mimetype()
                if mimetype is not None and mimetype.startswith('image'):
                    photos.append(f)
                elif filename == 'DESCRIPTION':
                    with open(f.get_realpath(), 'r') as f:
                        description = f.read()

        self.ctx.template_vars['dirs'] = dirs
        self.ctx.template_vars['photos'] = photos
        self.ctx.template_vars['description'] = description
        self.ctx.res.body = self.ctx.render('list-gallery.html')

class DownloadThumbnailAction(ViewAction):
    DEFAULT_SIZE = 300

    def _get_size(self):
        try:
            size = int(self.ctx.req.str_GET['thumb_size'])
        except (KeyError,ValueError):
            size = self.DEFAULT_SIZE
        else:
            if size < 1 or size > 1000:
                size = self.DEFAULT_SIZE
        return size

    def get(self):
        size = self._get_size()
        thumbdir = os.path.join(self.ctx.storage.path, 'thumbnails')
        f = self.ctx.file

        # use a lossless format in doubt
        lossy = f.get_mimetype() == 'image/jpeg'
        thumbext = 'jpg' if lossy else 'png'
        thumbpath = os.path.join(thumbdir, str(size), '%s.%s' % (f.get_hash(), thumbext))
        f.data['thumbnail']['ext'] = thumbext
        f.save()

        mtime = int(os.path.getmtime(f.get_realpath()))
        if not os.path.exists(thumbpath) or mtime > int(os.path.getmtime(thumbpath)):
            if not os.path.isdir(os.path.dirname(thumbpath)):
                os.makedirs(os.path.dirname(thumbpath))
            with open(f.get_realpath(), 'rb') as fp:
                img = Image.open(fp)
                img.thumbnail((size, size), Image.BILINEAR)
                with open(thumbpath, 'wb') as tfp:
                    if lossy:
                        img.save(tfp, 'jpeg', quality=95)
                    else:
                        img.save(tfp, 'png')
            os.utime(thumbpath, (mtime, mtime))

        self.ctx.res = FileApp(thumbpath)
        self.ctx.res.cache_control(private=True, max_age=CACHE_CONTROL.ONE_HOUR)

class GalleryPlugin(Plugin):
    def init(self):
        self.register_web_view(
            View(object_type='directory', name='gallery'),
            ListGalleryAction)
        self.register_web_view(
            View(object_type='file', mimetype='image', name='thumbnail'),
            DownloadThumbnailAction)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""\
This is a python port of "Goose" orignialy licensed to Gravity.com
under one or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.

Python port was written by Xavier Grangier for Recrutae

Gravity.com licenses this file
to you 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.
"""
import six
import requests

class NetworkError(RuntimeError):
    def __init__(self, status_code, reason):
        self.reason = reason
        self.status_code = status_code

class NetworkFetcher(object):

    def __init__(self, config):
        self.config = config
        self._connection = requests.Session()
        self._connection.headers['User-agent'] = self.config.browser_user_agent

        self._url = None
        self.response = None
        self.headers = None

    def get_url(self):
        return self._url

    def fetch(self, url):
        # utf-8 encode unicode url
        if isinstance(url, six.text_type) and six.PY2:
            url = url.encode('utf-8')

        self.response = self._connection.get(url, timeout=self.config.http_timeout, headers=self.headers)
        if self.response.ok:
            self._url = self.response.url
            text = self.response.content
        else:
            self._url = None
            text = None
            if self.config.strict:
                raise NetworkError(self.response.status_code, self.response.reason)

        return text
<commit_msg>remove self from response<commit_after># -*- coding: utf-8 -*-
"""\
This is a python port of "Goose" orignialy licensed to Gravity.com
under one or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.

Python port was written by Xavier Grangier for Recrutae

Gravity.com licenses this file
to you 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.
"""
import six
import requests

class NetworkError(RuntimeError):
    def __init__(self, status_code, reason):
        self.reason = reason
        self.status_code = status_code

class NetworkFetcher(object):

    def __init__(self, config):
        self.config = config
        self._connection = requests.Session()
        self._connection.headers['User-agent'] = self.config.browser_user_agent

        self._url = None
        self.response = None
        self.headers = None

    def get_url(self):
        return self._url

    def fetch(self, url):
        # utf-8 encode unicode url
        if isinstance(url, six.text_type) and six.PY2:
            url = url.encode('utf-8')

        response = self._connection.get(url, timeout=self.config.http_timeout, headers=self.headers)
        if response.ok:
            self._url = response.url
            text = response.content
        else:
            self._url = None
            text = None
            if self.config.strict:
                raise NetworkError(response.status_code, response.reason)

        return text
<|endoftext|>


<commit_before># Twitter Toolbox for Python
# Copyright 2016 Hugo Hromic
#
# 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.

"""Twitter Public Streaming API module."""

import logging
import json
from tweepy import StreamListener, Stream
from .helpers import init_logger, read_config, get_oauth_api
from .helpers import ensure_at_least_one

# module logging
LOGGER = logging.getLogger(__name__)
init_logger(LOGGER)

class PassThroughStreamListener(StreamListener):
    """Stream Listener that passes incoming messages directly to a writer."""

    def __init__(self, writer, limit=0, **kwargs):
        super(self.__class__, self).__init__(**kwargs)
        self.writer = writer
        self.limit = limit
        self.num_written = 0

    def on_status(self, status):
        """Write an incoming Tweet to the writer."""
        self.writer.write("%s\n" % json.dumps(status._json, separators=(",", ":")))  # pylint: disable=protected-access
        self.num_written += 1
        if self.num_written == self.limit:
            return False
        return True

    def on_error(self, status_code):
        """Handle stream errors."""
        if status_code == 420:
            LOGGER.error("too many connection attempts, stopping stream")
            return False

def _get_stream(writer, config, limit=0):
    api = get_oauth_api(config)
    listener = PassThroughStreamListener(writer, limit=limit)
    return Stream(auth=api.auth, listener=listener)

def get_sample(writer):
    """Get hydrated Tweet-objects from the sample Streaming API endpoint."""
    LOGGER.info("get_sample() starting")

    # initialize a Streaming API object and run the endpoint
    config = read_config()
    limit = config.getint("sample", "limit")
    stream = _get_stream(writer, config, limit=limit)
    stream.sample()

    # finished
    LOGGER.info("get_sample() finished")

def get_filter(writer, follow=None, track=None, locations=None):
    """Get hydrated Tweet-objects from the filter Streaming API endpoint."""
    LOGGER.info("get_filter() starting")
    ensure_at_least_one(follow=follow, track=track, locations=locations)

    # initialize a Streaming API object and run the endpoint
    config = read_config()
    limit = config.getint("filter", "limit")
    stream = _get_stream(writer, config, limit=limit)
    stream.filter(follow=follow, track=track, locations=locations)

    # finished
    LOGGER.info("get_filter() finished")

def get_firehose(writer):
    """Get hydrated Tweet-objects from the firehose Streaming API endpoint."""
    LOGGER.info("get_firehose() starting")

    # initialize a Streaming API object and run the endpoint
    config = read_config()
    limit = config.getint("firehose", "limit")
    stream = _get_stream(writer, config, limit=limit)
    stream.firehose()

    # finished
    LOGGER.info("get_firehose() finished")
<commit_msg>tweepy needs strings for user ids here<commit_after># Twitter Toolbox for Python
# Copyright 2016 Hugo Hromic
#
# 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.

"""Twitter Public Streaming API module."""

import logging
import json
from tweepy import StreamListener, Stream
from .helpers import init_logger, read_config, get_oauth_api
from .helpers import ensure_at_least_one

# module logging
LOGGER = logging.getLogger(__name__)
init_logger(LOGGER)

class PassThroughStreamListener(StreamListener):
    """Stream Listener that passes incoming messages directly to a writer."""

    def __init__(self, writer, limit=0, **kwargs):
        super(self.__class__, self).__init__(**kwargs)
        self.writer = writer
        self.limit = limit
        self.num_written = 0

    def on_status(self, status):
        """Write an incoming Tweet to the writer."""
        self.writer.write("%s\n" % json.dumps(status._json, separators=(",", ":")))  # pylint: disable=protected-access
        self.num_written += 1
        if self.num_written == self.limit:
            return False
        return True

    def on_error(self, status_code):
        """Handle stream errors."""
        if status_code == 420:
            LOGGER.error("too many connection attempts, stopping stream")
            return False

def _get_stream(writer, config, limit=0):
    api = get_oauth_api(config)
    listener = PassThroughStreamListener(writer, limit=limit)
    return Stream(auth=api.auth, listener=listener)

def get_sample(writer):
    """Get hydrated Tweet-objects from the sample Streaming API endpoint."""
    LOGGER.info("get_sample() starting")

    # initialize a Streaming API object and run the endpoint
    config = read_config()
    limit = config.getint("sample", "limit")
    stream = _get_stream(writer, config, limit=limit)
    stream.sample()

    # finished
    LOGGER.info("get_sample() finished")

def get_filter(writer, follow=None, track=None, locations=None):
    """Get hydrated Tweet-objects from the filter Streaming API endpoint."""
    LOGGER.info("get_filter() starting")
    ensure_at_least_one(follow=follow, track=track, locations=locations)

    # initialize a Streaming API object and run the endpoint
    config = read_config()
    limit = config.getint("filter", "limit")
    stream = _get_stream(writer, config, limit=limit)
    stream.filter(follow=[str(f) for f in follow], track=track, locations=locations)

    # finished
    LOGGER.info("get_filter() finished")

def get_firehose(writer):
    """Get hydrated Tweet-objects from the firehose Streaming API endpoint."""
    LOGGER.info("get_firehose() starting")

    # initialize a Streaming API object and run the endpoint
    config = read_config()
    limit = config.getint("firehose", "limit")
    stream = _get_stream(writer, config, limit=limit)
    stream.firehose()

    # finished
    LOGGER.info("get_firehose() finished")
<|endoftext|>