<commit_before>import py
from prolog.interpreter.test.tool import assert_true, \
get_engine, assert_false, collect_all
from prolog.interpreter.continuation import Engine

loc = "../../prolog_modules/"
e = get_engine("""
:- use_module('%s%s').
""" % (loc, "dcg"))

def test_a_to_b():
    assert_true("""
    trans((a --> b), X),
    X = (a(X1, X2) :- b(X1, X2)).
    """, e)

def test_a_to_b_c():
    assert_true("""
    trans((a --> b, c), X),
    X = (a(X1, X3) :- b(X1, X2), c(X2, X3)).
    """, e)

def test_single_terminal():
    assert_true("""
    trans((a --> [b]), X),
    X = (a([b|X1], X1) :- true).
    """, e)

def test_terminal_nonterminal():
    assert_true("""
    trans((a --> [c], b), X),
    X = (a([c|Y1], Y2) :- b(Y1, Y2)).
    """, e)

def test_nonterminal_terminal():
    assert_true("""
    trans((a --> b, [c]), X),
    X = (a(X1, X3) :- b(X1, X2), X2=[c|X3]).
    """, e)

def test_nonterminal_terminal_terminal():
    assert_true("""
    trans((a --> b, [c], [d]), X),
    X = (a(X1, X3) :- b(X1, X2), X2=[c, d|X3]).
    """, e)

def test_terminal_terminal():
    assert_true("""
    trans((a --> [b], [c]), X),
    X = (a([b, c|X], X) :- true).
    """, e)

def test_nonterminal_param_1():
    assert_true("""
    trans((a(X) --> b), R),
    R = (a(X, X1, X2) :- b(X1, X2)).
    """, e)

def test_nonterminal_param_2():
    assert_true("""
    trans((a(X, Y) --> b(Y)), R),
    R = (a(X, Y, X1, X2) :- b(Y, X1, X2)).
    """, e)

def test_curly():
    assert_true("""
    trans((a --> {X=1}), R),
    R = (a(X1, X2) :- X=1, X2=X1).
    """, e)

def test_nonterminal_curly():
    assert_true("""
    trans((a --> b, {X=1}), R),
    R = (a(X1, X2) :- b(X1, X3), X=1, X2=X3).
    """, e)

def test_nonterminal_nonterminal_curly():
    assert_true("""
    trans((a --> b, c, {X=1}), R),
    R = (a(X1, X2) :- b(X1, X3), c(X3, X4), X=1, X2=X4).
    """, e)

def test_curly_nonterminal():
    assert_true("""
    trans((a --> {X=1}, b), R),
    R = (a(X1, X2) :- X=1, b(X1, X2)).
    """, e)

def test_curly_multiple_nonterminal():
    assert_true("""
    trans((a --> {X=1, Y=2, Z=3}, b), R),
    R = (a(X1, X2) :- (X=1, Y=2, Z=3), b(X1, X2)).
    """, e)

def test_curly_nonterminal_curly_nonterminal():
    assert_true("""
    trans((a --> {X=1}, b, {Y=2}, c), R),
    R = (a(X1, X2) :- X=1, b(X1, X3), Y=2, c(X3, X2)).
    """, e)

def test_curly_terminal():
    assert_true("""
    trans((a --> {X=1}, [b]), R),
    R = (a(X1, X2) :- X=1, X1=[b|X2]).
    """, e)

def test_dcg_integration_1():
    assert_true("""
    trans((a --> {X=1, Y=2}, [b, c], b(Y)), R),
    R = (a(X1, X2) :- (X=1, Y=2), X1=[b, c|X4], b(Y, X4, X2)).
    """, e)

def test_curly_multiple_1():
    assert_true("""
    trans((a --> {X=1, Y=2}), R),
    R = (a(X1, X2) :- X=1, Y=2, X1=X2).
    """, e)

def test_curly_multiple_2():
    assert_true("""
    trans((a --> {X=1, Y=2, Z=3}), R),
    R = (a(X1, X2) :- X=1, Y=2, Z=3, X1=X2).
    """, e)

# DCG integration tests, i. e. check translated rules
# ___________________________________________________

def make_engine():
    e = get_engine("""
    :- use_module('%s%s').
    """ % (loc, "dcg"))
    return e

def test_integration_1():
    assert_true("""
    trans((a1 --> [b], [c]), R1),
    assert(R1),
    a1([b, c], []).
    """, e)

def test_integration_2():
    assert_true("""
    trans((a2 --> b2, c2, [d]), R1),
    trans((b2 --> [b]), R2),
    trans((c2 --> [c]), R3),
    assert(R1), assert(R2), assert(R3),
    a2([b, c, d], []).
    """, e)

def test_integration_3():
    assert_true("""
    trans((s --> np, vp), R1), assert(R1),
    trans((np --> det, n), R2), assert(R2),
    trans((vp --> tv, np), R3), assert(R3),
    trans((vp --> v), R4), assert(R4),
    trans((det --> [the]), R5), assert(R5),
    trans((det --> [a]), R6), assert(R6),
    trans((det --> [every]), R7), assert(R7),
    trans((n --> [man]), R8), assert(R8),
    trans((n --> [woman]), R9), assert(R9),
    trans((n --> [park]), R10), assert(R10),
    trans((tv --> [loves]), R11), assert(R11),
    trans((tv --> [likes]), R12), assert(R12),
    trans((n --> [park]), R13), assert(R13),
    s([a, man, loves, the, woman], []),
    s([every, woman, likes, a, park], []).
    """, e)
    assert_false("""
    s([every, every, woman, likes, a, man], []).
    """, e)
    assert_true("s([the, woman, loves, the, man], []).", e)

def test_integration_4():
    assert_true("""
    trans((sb --> ab, [or], sb), R1), assert(R1),
    trans((sb --> ab), R2), assert(R2),
    trans((ab --> bb, [and], sb), R3), assert(R3),
    trans((ab --> bb), R4), assert(R4),
    trans((bb --> [true]), R5), assert(R5),
    trans((bb --> [false]), R6), assert(R6),
    trans((bb --> ['('], sb, [')']), R7), assert(R7),
    trans((bb --> [not, '('], sb, [')']), R8), assert(R8).
    """, e)
    assert_true("sb([true, or, false], []).", e)
    assert_true("sb([false, and, '(', false, ')'], []).", e)
    assert_true("sb([not, '(', true, ')'], []).", e)
    assert_true("sb([true, and, false, or, true, and, not, '(', false, or, true, ')'], []).", e)
    assert_true("sb(['(', false, ')'], []).", e)

    assert_false("sb([true, true], []).", e)
    assert_false("sb([true, '(', false], []).", e)
    assert_false("sb([true, or, true, and, not, false], []).", e)
    assert_false("sb([not, '(', false, ')', or], []).", e)

def test_integration_5():
    assert_true("""
    trans((expr --> term, addterm), R1), assert(R1),
    trans((addterm --> []), R2), assert(R2),
    trans((addterm --> [+], expr), R3), assert(R3),
    trans((term --> factor, multfactor), R4), assert(R4),
    trans((multfactor --> []), R5), assert(R5),
    trans((multfactor --> [*], term), R6), assert(R6),
    trans((factor --> [I], {integer(I)}), R7), assert(R7),
    trans((factor --> ['('], expr, [')']), R8), assert(R8).
    """, e)

    assert_true("expr([1, '+', 2], []).", e)
    assert_true("expr([1, '*', '(', 2, '+', 3, ')'], []).", e)
    assert_true("expr([1, '*', '(', 2, '+', '(', 3, ')', ')'], []).", e)
    assert_true("expr([1000], []).", e)

    assert_false("expr([100, 1000], []).", e)
    assert_false("expr([a], []).", e)
    assert_false("expr([1, '+', '*', '2'], []).", e)
<commit_msg>altered integration testswq<commit_after>import py
from prolog.interpreter.test.tool import assert_true, \
get_engine, assert_false, collect_all
from prolog.interpreter.continuation import Engine

loc = "../../prolog_modules/"
e = get_engine("""
:- use_module('%s%s').
""" % (loc, "dcg"))

def test_a_to_b():
    assert_true("""
    trans((a --> b), X),
    X = (a(X1, X2) :- b(X1, X2)).
    """, e)

def test_a_to_b_c():
    assert_true("""
    trans((a --> b, c), X),
    X = (a(X1, X3) :- b(X1, X2), c(X2, X3)).
    """, e)

def test_single_terminal():
    assert_true("""
    trans((a --> [b]), X),
    X = (a([b|X1], X1) :- true).
    """, e)

def test_terminal_nonterminal():
    assert_true("""
    trans((a --> [c], b), X),
    X = (a([c|Y1], Y2) :- b(Y1, Y2)).
    """, e)

def test_nonterminal_terminal():
    assert_true("""
    trans((a --> b, [c]), X),
    X = (a(X1, X3) :- b(X1, X2), X2=[c|X3]).
    """, e)

def test_nonterminal_terminal_terminal():
    assert_true("""
    trans((a --> b, [c], [d]), X),
    X = (a(X1, X3) :- b(X1, X2), X2=[c, d|X3]).
    """, e)

def test_terminal_terminal():
    assert_true("""
    trans((a --> [b], [c]), X),
    X = (a([b, c|X], X) :- true).
    """, e)

def test_nonterminal_param_1():
    assert_true("""
    trans((a(X) --> b), R),
    R = (a(X, X1, X2) :- b(X1, X2)).
    """, e)

def test_nonterminal_param_2():
    assert_true("""
    trans((a(X, Y) --> b(Y)), R),
    R = (a(X, Y, X1, X2) :- b(Y, X1, X2)).
    """, e)

def test_curly():
    assert_true("""
    trans((a --> {X=1}), R),
    R = (a(X1, X2) :- X=1, X2=X1).
    """, e)

def test_nonterminal_curly():
    assert_true("""
    trans((a --> b, {X=1}), R),
    R = (a(X1, X2) :- b(X1, X3), X=1, X2=X3).
    """, e)

def test_nonterminal_nonterminal_curly():
    assert_true("""
    trans((a --> b, c, {X=1}), R),
    R = (a(X1, X2) :- b(X1, X3), c(X3, X4), X=1, X2=X4).
    """, e)

def test_curly_nonterminal():
    assert_true("""
    trans((a --> {X=1}, b), R),
    R = (a(X1, X2) :- X=1, b(X1, X2)).
    """, e)

def test_curly_multiple_nonterminal():
    assert_true("""
    trans((a --> {X=1, Y=2, Z=3}, b), R),
    R = (a(X1, X2) :- (X=1, Y=2, Z=3), b(X1, X2)).
    """, e)

def test_curly_nonterminal_curly_nonterminal():
    assert_true("""
    trans((a --> {X=1}, b, {Y=2}, c), R),
    R = (a(X1, X2) :- X=1, b(X1, X3), Y=2, c(X3, X2)).
    """, e)

def test_curly_terminal():
    assert_true("""
    trans((a --> {X=1}, [b]), R),
    R = (a(X1, X2) :- X=1, X1=[b|X2]).
    """, e)

def test_dcg_integration_1():
    assert_true("""
    trans((a --> {X=1, Y=2}, [b, c], b(Y)), R),
    R = (a(X1, X2) :- (X=1, Y=2), X1=[b, c|X4], b(Y, X4, X2)).
    """, e)

def test_curly_multiple_1():
    assert_true("""
    trans((a --> {X=1, Y=2}), R),
    R = (a(X1, X2) :- X=1, Y=2, X1=X2).
    """, e)

def test_curly_multiple_2():
    assert_true("""
    trans((a --> {X=1, Y=2, Z=3}), R),
    R = (a(X1, X2) :- X=1, Y=2, Z=3, X1=X2).
    """, e)

# DCG integration tests, i. e. check translated rules
# ___________________________________________________

def test_integration_1():
    assert_true("""
    trans((a1 --> [b], [c]), R1),
    assert(R1),
    a1([b, c], []).
    """, e)

def test_integration_2():
    assert_true("""
    trans((a2 --> b2, c2, [d]), R1),
    trans((b2 --> [b]), R2),
    trans((c2 --> [c]), R3),
    assert(R1), assert(R2), assert(R3),
    a2([b, c, d], []).
    """, e)

def test_integration_3():
    assert_true("""
    trans((x --> blub, zip), R1), assert(R1),
    trans((blub --> pre, tgz), R2), assert(R2),
    trans((zip --> bla, blub), R3), assert(R3),
    trans((zip --> v), R4), assert(R4),
    trans((pre --> [the]), R5), assert(R5),
    trans((pre --> [a]), R6), assert(R6),
    trans((pre --> [that]), R7), assert(R7),
    trans((tgz --> [bug]), R8), assert(R8),
    trans((tgz --> [witch]), R9), assert(R9),
    trans((tgz --> [besom]), R10), assert(R10),
    trans((tgz --> [pipe]), R14), assert(R14),
    trans((bla --> [hates]), R11), assert(R11),
    trans((bla --> [fools]), R12), assert(R12),
    trans((n --> [besom]), R13), assert(R13),
    x([a, bug, hates, the, witch], []),
    x([that, witch, fools, a, besom], []).
    """, e)
    assert_false("""
    x([that, that, witch, fools, a, bug], []).
    """, e)
    assert_true("x([the, witch, hates, the, bug], []).", e)

def test_integration_4():
    assert_true("""
    trans((sb --> ab, [or], sb), R1), assert(R1),
    trans((sb --> ab), R2), assert(R2),
    trans((ab --> bb, [and], sb), R3), assert(R3),
    trans((ab --> bb), R4), assert(R4),
    trans((bb --> [true]), R5), assert(R5),
    trans((bb --> [false]), R6), assert(R6),
    trans((bb --> ['('], sb, [')']), R7), assert(R7),
    trans((bb --> [not, '('], sb, [')']), R8), assert(R8).
    """, e)
    assert_true("sb([true, or, false], []).", e)
    assert_true("sb([false, and, '(', false, ')'], []).", e)
    assert_true("sb([not, '(', true, ')'], []).", e)
    assert_true("sb([true, and, false, or, true, and, not, '(', false, or, true, ')'], []).", e)
    assert_true("sb(['(', false, ')'], []).", e)

    assert_false("sb([true, true], []).", e)
    assert_false("sb([true, '(', false], []).", e)
    assert_false("sb([true, or, true, and, not, false], []).", e)
    assert_false("sb([not, '(', false, ')', or], []).", e)

def test_integration_5():
    assert_true("""
    trans((ex --> tm, at), R1), assert(R1),
    trans((at --> []), R2), assert(R2),
    trans((at --> [+], ex), R3), assert(R3),
    trans((tm --> fcr, mf), R4), assert(R4),
    trans((mf --> []), R5), assert(R5),
    trans((mf --> [*], tm), R6), assert(R6),
    trans((fcr --> [I], {integer(I)}), R7), assert(R7),
    trans((fcr --> ['('], ex, [')']), R8), assert(R8).
    """, e)

    assert_true("ex([1, '+', 2], []).", e)
    assert_true("ex([1, '*', '(', 2, '+', 3, ')'], []).", e)
    assert_true("ex([1, '*', '(', 2, '+', '(', 3, ')', ')'], []).", e)
    assert_true("ex([1000], []).", e)

    assert_false("ex([100, 1000], []).", e)
    assert_false("ex([a], []).", e)
    assert_false("ex([1, '+', '*', '2'], []).", e)
<|endoftext|>


<commit_before>import os
import tempfile
import logging
from pkg_resources import resource_filename
import contextlib
import shutil
import mdtraj as md
from mdtraj.utils import enter_temp_directory

try:
    from subprocess import getoutput  # If python 3
except ImportError:
    from commands import getoutput  # If python 2

import simtk.openmm
from simtk.openmm import app
import simtk.unit as units
from distutils.spawn import find_executable

from gaff2xml import amber_parser, system_checker

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG, format="LOG: %(message)s")

def find_gaff_dat():
    AMBERHOME = None

    try:
        AMBERHOME = os.environ['AMBERHOME']
    except KeyError:
        pass

    if AMBERHOME is None:
        full_path = find_executable("parmchk2")
        try:
            AMBERHOME = os.path.split(full_path)[0]
            AMBERHOME = os.path.join(AMBERHOME, "../")
        except:
            raise(ValueError("Cannot find AMBER GAFF"))

    if AMBERHOME is None:
        raise(ValueError("Cannot find AMBER GAFF"))

    return os.path.join(AMBERHOME, 'dat', 'leap', 'parm', 'gaff.dat')

GAFF_DAT_FILENAME = find_gaff_dat()

def parse_ligand_filename(filename):
    """Split ligand filename into name and extension.  "./ligand.mol2" -> ("ligand", ".mol2")"""
    name, ext = os.path.splitext(os.path.split(filename)[1])
    return name, ext

def run_antechamber(molecule_name, input_filename, charge_method="bcc", net_charge=None):
    """Run AmberTools antechamber and parmchk2 to create GAFF mol2 and frcmod files.

    Parameters
    ----------
    molecule_name : str
        Name of the molecule to be parameterized, will be used in output filenames.
    ligand_filename : str
        The molecule to be parameterized.  Must be tripos mol2 format.
    charge_method : str, optional
        If not None, the charge method string will be passed to Antechamber.
    net_charge : int, optional
        If not None, net charge of the molecule to be parameterized.
        If None, Antechamber sums up partial charges from the input file.

    Returns
    -------
    gaff_mol2_filename : str
        GAFF format mol2 filename produced by antechamber
    frcmod_filename : str
        Amber frcmod file produced by prmchk
    """

    ext = parse_ligand_filename(input_filename)[1]

    filetype = ext[1:]
    if filetype != "mol2":
        raise(ValueError("Must input mol2 filename"))

    gaff_mol2_filename = molecule_name + '.gaff.mol2'
    frcmod_filename = molecule_name + '.frcmod'

    cmd = "antechamber -i %s -fi mol2 -o %s -fo mol2 -s 2" % (input_filename, gaff_mol2_filename)
    if charge_method is not None:
        cmd += ' -c %s' % charge_method

    if net_charge is not None:
        cmd += ' -nc %d' % net_charge

    logger.debug(cmd)

    output = getoutput(cmd)
    logger.debug(output)

    cmd = "parmchk2 -i %s -f mol2 -o %s" % (gaff_mol2_filename, frcmod_filename)
    logger.debug(cmd)

    output = getoutput(cmd)
    logger.debug(output)

    return gaff_mol2_filename, frcmod_filename

def convert_molecule(in_filename, out_filename):
    """Use openbabel to convert filenames.  May not work for all file formats!"""

    molecule_name, ext_in = parse_ligand_filename(in_filename)
    molecule_name, ext_out = parse_ligand_filename(out_filename)

    cmd = "obabel -i %s %s -o %s -O %s" % (ext_in[1:], in_filename, ext_out[1:], out_filename)
    print(cmd)
    output = getoutput(cmd)
    logger.debug(output)

def run_tleap(molecule_name, gaff_mol2_filename, frcmod_filename):
    """Run AmberTools tleap to create simulation files for AMBER

    Parameters
    ----------
    gaff_mol2_filename : str
        GAFF format mol2 filename produced by antechamber
    frcmod_filename : str
        Amber frcmod file produced by prmchk

    Returns
    -------
    gaff_mol2_filename : str
        GAFF format mol2 filename produced by antechamber
    frcmod_filename : str
        Amber frcmod file produced by prmchk
    """

    prmtop_filename = "%s.prmtop" % molecule_name
    inpcrd_filename = "%s.inpcrd" % molecule_name

    tleap_input = """
source leaprc.ff99SB
source leaprc.gaff
LIG = loadmol2 %s
check LIG
loadamberparams %s
saveamberparm LIG %s %s
quit

""" % (gaff_mol2_filename, frcmod_filename, prmtop_filename, inpcrd_filename)

    file_handle = tempfile.NamedTemporaryFile('w')  # FYI Py3K defaults to 'wb' mode, which won't work here.
    file_handle.writelines(tleap_input)
    file_handle.flush()

    cmd = "tleap -f %s " % file_handle.name
    logger.debug(cmd)

    output = getoutput(cmd)
    logger.debug(output)

    file_handle.close()

    return prmtop_filename, inpcrd_filename

def molecule_to_mol2(molecule, tripos_mol2_filename=None):
    """Convert OE molecule to tripos mol2 file.

    Parameters
    ----------
    molecule : openeye.oechem.OEGraphMol
        The molecule to be converted.

    Returns
    -------
    tripos_mol2_filename : str
        Filename of output tripos mol2 file

    """

    try:
        import openeye.oechem
    except ImportError:
        raise(ImportError("Must install OpenEye tools to process OpenEye MOL2 files."))

    # Get molecule name.
    molecule_name = molecule.GetTitle()
    logger.debug(molecule_name)

    # Write molecule as Tripos mol2.
    if tripos_mol2_filename is None:
        tripos_mol2_filename = molecule_name + '.tripos.mol2'

    ofs = openeye.oechem.oemolostream(tripos_mol2_filename)
    ofs.SetFormat(openeye.oechem.OEFormat_MOL2H)
    openeye.oechem.OEWriteMolecule(ofs, molecule)
    ofs.close()

    # Replace <0> substructure names with valid text.
    infile = open(tripos_mol2_filename, 'r')
    lines = infile.readlines()
    infile.close()
    newlines = [line.replace('<0>', 'MOL') for line in lines]
    outfile = open(tripos_mol2_filename, 'w')
    outfile.writelines(newlines)
    outfile.close()

    return molecule_name, tripos_mol2_filename

def create_ffxml_file(gaff_mol2_filenames, frcmod_filenames, ffxml_filename=None, override_mol2_residue_name=None):
    """Process a gaff mol2 file and frcmod file using the XML conversion and write to an XML file.

    Parameters
    ----------
    gaff_mol2_filenames : list of str
        The names of the gaff mol2 files
    frcmod_filenames : str
        The names of the gaff frcmod files
    ffxml_filename : str, optional, default=None
        Optional name of output ffxml file to generate.  If None, no file
        will be generated.
    override_mol2_residue_name : str, default=None
            If given, use this name to override mol2 residue names.

    Returns
    -------
    ffxml_stringio : str
        StringIO representation of ffxml file containing residue entries for each molecule.

    """

    # Generate ffxml file.
    parser = amber_parser.AmberParser(override_mol2_residue_name=override_mol2_residue_name)

    filenames = [GAFF_DAT_FILENAME]
    filenames.extend([filename for filename in gaff_mol2_filenames])
    filenames.extend([filename for filename in frcmod_filenames])

    parser.parse_filenames(filenames)

    ffxml_stream = parser.generate_xml()

    if ffxml_filename is not None:
        outfile = open(ffxml_filename, 'w')
        outfile.write(ffxml_stream.read())
        outfile.close()
        ffxml_stream.seek(0)

    return ffxml_stream

def create_ffxml_simulation(molecule_name, gaff_mol2_filename, frcmod_filename):
    """Process a gaff mol2 file and frcmod file using the XML conversion, returning an OpenMM simulation.

    Parameters
    ----------
    molecule_name : str
        The name of the molecule
    gaff_mol2_filename : str
        The name of the gaff mol2 file
    frcmod_filename : str
        The name of the gaff frcmod file

    Returns
    -------
    simulation : openmm.app.Simulation
        A functional simulation object for simulating your molecule
    """

    # Generate ffxml file.
    parser = amber_parser.AmberParser()
    parser.parse_filenames([GAFF_DAT_FILENAME, gaff_mol2_filename, frcmod_filename])

    ffxml_filename = molecule_name + '.ffxml'
    create_ffxml_file([gaff_mol2_filename], [frcmod_filename], ffxml_filename)

    traj = md.load(gaff_mol2_filename)  # Read mol2 file.
    positions = traj.openmm_positions(0)  # Extract OpenMM-united positions of first (and only) trajectory frame
    topology = traj.top.to_openmm()

    # Create System object.
    forcefield = app.ForceField(ffxml_filename)
    system = forcefield.createSystem(topology, nonbondedMethod=app.NoCutoff, constraints=None, implicitSolvent=None)

    # Create integrator.
    timestep = 1.0 * units.femtoseconds
    integrator = simtk.openmm.VerletIntegrator(timestep)

    # Create simulation.
    platform = simtk.openmm.Platform.getPlatformByName("Reference")
    simulation = app.Simulation(topology, system, integrator, platform=platform)
    simulation.context.setPositions(positions)

    return simulation

def create_leap_simulation(molecule_name, gaff_mol2_filename, frcmod_filename):
    """Create an OpenMM simulation using a Gaff mol2 file and frcmod file.

    Parameters
    ----------
    molecule_name : str
        Name of the molecule
    gaff_mol2_filename : str
        Filename of input (GAFF!) mol2 file
    frcmod_filename : str
        Use this frcmod filename

    """

    # Parameterize system with LEaP.
    (prmtop_filename, inpcrd_filename) = run_tleap(molecule_name, gaff_mol2_filename, frcmod_filename)

    # Create System object.
    prmtop = app.AmberPrmtopFile(prmtop_filename)
    topology = prmtop.topology
    system = prmtop.createSystem(nonbondedMethod=app.NoCutoff, constraints=None, implicitSolvent=None)

    # Read positions.
    inpcrd = app.AmberInpcrdFile(inpcrd_filename)
    positions = inpcrd.getPositions()

    # Create integrator.
    timestep = 1.0 * units.femtoseconds
    integrator = simtk.openmm.VerletIntegrator(timestep)

    platform = simtk.openmm.Platform.getPlatformByName("Reference")
    simulation = app.Simulation(topology, system, integrator, platform=platform)
    simulation.context.setPositions(positions)

    return simulation

def test_molecule(molecule_name, tripos_mol2_filename, charge_method=None):
    """Create a GAFF molecule via LEAP and ffXML and compare force terms.

    Parameters
    ----------
    molecule_name : str
        Name of the molecule
    tripos_mol2_filename : str
        Filename of input mol2 file
    charge_method : str, default=None
        If None, use charges in existing MOL2.  Otherwise, use a charge
        model when running antechamber.
    """

    # Generate GAFF parameters.
    (gaff_mol2_filename, frcmod_filename) = run_antechamber(molecule_name, tripos_mol2_filename, charge_method=charge_method)

    # Create simulations.
    simulation_ffxml = create_ffxml_simulation(molecule_name, gaff_mol2_filename, frcmod_filename)
    simulation_leap  = create_leap_simulation(molecule_name, gaff_mol2_filename, frcmod_filename)

    # Compare simulations.
    syscheck = system_checker.SystemChecker(simulation_ffxml, simulation_leap)
    syscheck.check_force_parameters()

    groups0, groups1 = syscheck.check_energy_groups()
    energy0, energy1 = syscheck.check_energies()

def get_data_filename(relative_path):
    """Get the full path to one of the reference files shipped for testing

    In the source distribution, these files are in ``gaff2xml/chemicals/*/``,
    but on installation, they're moved to somewhere in the user's python
    site-packages directory.

    Parameters
    ----------
    name : str
        Name of the file to load (with respect to the gaff2xml folder).

    """

    fn = resource_filename('gaff2xml', relative_path)

    if not os.path.exists(fn):
        raise ValueError("Sorry! %s does not exist. If you just added it, you'll have to re-install" % fn)

    return fn

def smiles_to_mdtraj_ffxml(smiles_strings, base_molecule_name="lig"):
    """Generate an MDTraj object from a smiles string.

    Parameters
    ----------
    smiles_strings : list(str)
        Smiles strings to create molecules for
    base_molecule_name : str, optional, default='lig'
        Base name of molecule to use inside parameter files.

    Returns
    -------
    traj : mdtraj.Trajectory
        MDTraj object for molecule
    ffxml : StringIO
        StringIO representation of ffxml file.

    Notes
    -----
    ffxml can be directly input to OpenMM e.g.
    `forcefield = app.ForceField(ffxml)`
    """
    try:
        from rdkit import Chem
        from rdkit.Chem import AllChem
    except ImportError:
        raise(ImportError("Must install rdkit to use smiles conversion."))

    gaff_mol2_filenames = []
    frcmod_filenames = []
    trajectories = []
    for k, smiles_string in enumerate(smiles_strings):
        molecule_name = "%s-%d" % (base_molecule_name, k)
        m = Chem.MolFromSmiles(smiles_string)
        m = Chem.AddHs(m)
        AllChem.EmbedMolecule(m)
        AllChem.UFFOptimizeMolecule(m)

        pdb_filename = tempfile.mktemp(suffix=".pdb")
        Chem.MolToPDBFile(m, pdb_filename)

        mol2_filename = tempfile.mktemp(suffix=".mol2")

        convert_molecule(pdb_filename, mol2_filename)  # This is necessary because PDB double bonds are not handled by antechamber...
        print(mol2_filename)

        gaff_mol2_filename, frcmod_filename = run_antechamber(molecule_name, mol2_filename)
        traj = md.load(gaff_mol2_filename)
        print(gaff_mol2_filename)
        print(traj)

        for atom in traj.top.atoms:
            atom.residue.name = molecule_name

        gaff_mol2_filenames.append(gaff_mol2_filename)
        frcmod_filenames.append(frcmod_filename)
        trajectories.append(traj)

    ffxml = create_ffxml_file(gaff_mol2_filenames, frcmod_filenames, override_mol2_residue_name=molecule_name)

    return trajectories, ffxml

<commit_msg>Fixed incorrect return argument documentation in run_tleap<commit_after>import os
import tempfile
import logging
from pkg_resources import resource_filename
import contextlib
import shutil
import mdtraj as md
from mdtraj.utils import enter_temp_directory

try:
    from subprocess import getoutput  # If python 3
except ImportError:
    from commands import getoutput  # If python 2

import simtk.openmm
from simtk.openmm import app
import simtk.unit as units
from distutils.spawn import find_executable

from gaff2xml import amber_parser, system_checker

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG, format="LOG: %(message)s")

def find_gaff_dat():
    AMBERHOME = None

    try:
        AMBERHOME = os.environ['AMBERHOME']
    except KeyError:
        pass

    if AMBERHOME is None:
        full_path = find_executable("parmchk2")
        try:
            AMBERHOME = os.path.split(full_path)[0]
            AMBERHOME = os.path.join(AMBERHOME, "../")
        except:
            raise(ValueError("Cannot find AMBER GAFF"))

    if AMBERHOME is None:
        raise(ValueError("Cannot find AMBER GAFF"))

    return os.path.join(AMBERHOME, 'dat', 'leap', 'parm', 'gaff.dat')

GAFF_DAT_FILENAME = find_gaff_dat()

def parse_ligand_filename(filename):
    """Split ligand filename into name and extension.  "./ligand.mol2" -> ("ligand", ".mol2")"""
    name, ext = os.path.splitext(os.path.split(filename)[1])
    return name, ext

def run_antechamber(molecule_name, input_filename, charge_method="bcc", net_charge=None):
    """Run AmberTools antechamber and parmchk2 to create GAFF mol2 and frcmod files.

    Parameters
    ----------
    molecule_name : str
        Name of the molecule to be parameterized, will be used in output filenames.
    ligand_filename : str
        The molecule to be parameterized.  Must be tripos mol2 format.
    charge_method : str, optional
        If not None, the charge method string will be passed to Antechamber.
    net_charge : int, optional
        If not None, net charge of the molecule to be parameterized.
        If None, Antechamber sums up partial charges from the input file.

    Returns
    -------
    gaff_mol2_filename : str
        GAFF format mol2 filename produced by antechamber
    frcmod_filename : str
        Amber frcmod file produced by prmchk
    """

    ext = parse_ligand_filename(input_filename)[1]

    filetype = ext[1:]
    if filetype != "mol2":
        raise(ValueError("Must input mol2 filename"))

    gaff_mol2_filename = molecule_name + '.gaff.mol2'
    frcmod_filename = molecule_name + '.frcmod'

    cmd = "antechamber -i %s -fi mol2 -o %s -fo mol2 -s 2" % (input_filename, gaff_mol2_filename)
    if charge_method is not None:
        cmd += ' -c %s' % charge_method

    if net_charge is not None:
        cmd += ' -nc %d' % net_charge

    logger.debug(cmd)

    output = getoutput(cmd)
    logger.debug(output)

    cmd = "parmchk2 -i %s -f mol2 -o %s" % (gaff_mol2_filename, frcmod_filename)
    logger.debug(cmd)

    output = getoutput(cmd)
    logger.debug(output)

    return gaff_mol2_filename, frcmod_filename

def convert_molecule(in_filename, out_filename):
    """Use openbabel to convert filenames.  May not work for all file formats!"""

    molecule_name, ext_in = parse_ligand_filename(in_filename)
    molecule_name, ext_out = parse_ligand_filename(out_filename)

    cmd = "obabel -i %s %s -o %s -O %s" % (ext_in[1:], in_filename, ext_out[1:], out_filename)
    print(cmd)
    output = getoutput(cmd)
    logger.debug(output)

def run_tleap(molecule_name, gaff_mol2_filename, frcmod_filename):
    """Run AmberTools tleap to create simulation files for AMBER

    Parameters
    ----------
    gaff_mol2_filename : str
        GAFF format mol2 filename produced by antechamber
    frcmod_filename : str
        Amber frcmod file produced by prmchk

    Returns
    -------
    prmtop_filename : str
        Amber prmtop file produced by tleap
    inpcrd_filename : str
        Amber inpcrd file produced by tleap
    """

    prmtop_filename = "%s.prmtop" % molecule_name
    inpcrd_filename = "%s.inpcrd" % molecule_name

    tleap_input = """
source leaprc.ff99SB
source leaprc.gaff
LIG = loadmol2 %s
check LIG
loadamberparams %s
saveamberparm LIG %s %s
quit

""" % (gaff_mol2_filename, frcmod_filename, prmtop_filename, inpcrd_filename)

    file_handle = tempfile.NamedTemporaryFile('w')  # FYI Py3K defaults to 'wb' mode, which won't work here.
    file_handle.writelines(tleap_input)
    file_handle.flush()

    cmd = "tleap -f %s " % file_handle.name
    logger.debug(cmd)

    output = getoutput(cmd)
    logger.debug(output)

    file_handle.close()

    return prmtop_filename, inpcrd_filename

def molecule_to_mol2(molecule, tripos_mol2_filename=None):
    """Convert OE molecule to tripos mol2 file.

    Parameters
    ----------
    molecule : openeye.oechem.OEGraphMol
        The molecule to be converted.

    Returns
    -------
    tripos_mol2_filename : str
        Filename of output tripos mol2 file

    """

    try:
        import openeye.oechem
    except ImportError:
        raise(ImportError("Must install OpenEye tools to process OpenEye MOL2 files."))

    # Get molecule name.
    molecule_name = molecule.GetTitle()
    logger.debug(molecule_name)

    # Write molecule as Tripos mol2.
    if tripos_mol2_filename is None:
        tripos_mol2_filename = molecule_name + '.tripos.mol2'

    ofs = openeye.oechem.oemolostream(tripos_mol2_filename)
    ofs.SetFormat(openeye.oechem.OEFormat_MOL2H)
    openeye.oechem.OEWriteMolecule(ofs, molecule)
    ofs.close()

    # Replace <0> substructure names with valid text.
    infile = open(tripos_mol2_filename, 'r')
    lines = infile.readlines()
    infile.close()
    newlines = [line.replace('<0>', 'MOL') for line in lines]
    outfile = open(tripos_mol2_filename, 'w')
    outfile.writelines(newlines)
    outfile.close()

    return molecule_name, tripos_mol2_filename

def create_ffxml_file(gaff_mol2_filenames, frcmod_filenames, ffxml_filename=None, override_mol2_residue_name=None):
    """Process a gaff mol2 file and frcmod file using the XML conversion and write to an XML file.

    Parameters
    ----------
    gaff_mol2_filenames : list of str
        The names of the gaff mol2 files
    frcmod_filenames : str
        The names of the gaff frcmod files
    ffxml_filename : str, optional, default=None
        Optional name of output ffxml file to generate.  If None, no file
        will be generated.
    override_mol2_residue_name : str, default=None
            If given, use this name to override mol2 residue names.

    Returns
    -------
    ffxml_stringio : str
        StringIO representation of ffxml file containing residue entries for each molecule.

    """

    # Generate ffxml file.
    parser = amber_parser.AmberParser(override_mol2_residue_name=override_mol2_residue_name)

    filenames = [GAFF_DAT_FILENAME]
    filenames.extend([filename for filename in gaff_mol2_filenames])
    filenames.extend([filename for filename in frcmod_filenames])

    parser.parse_filenames(filenames)

    ffxml_stream = parser.generate_xml()

    if ffxml_filename is not None:
        outfile = open(ffxml_filename, 'w')
        outfile.write(ffxml_stream.read())
        outfile.close()
        ffxml_stream.seek(0)

    return ffxml_stream

def create_ffxml_simulation(molecule_name, gaff_mol2_filename, frcmod_filename):
    """Process a gaff mol2 file and frcmod file using the XML conversion, returning an OpenMM simulation.

    Parameters
    ----------
    molecule_name : str
        The name of the molecule
    gaff_mol2_filename : str
        The name of the gaff mol2 file
    frcmod_filename : str
        The name of the gaff frcmod file

    Returns
    -------
    simulation : openmm.app.Simulation
        A functional simulation object for simulating your molecule
    """

    # Generate ffxml file.
    parser = amber_parser.AmberParser()
    parser.parse_filenames([GAFF_DAT_FILENAME, gaff_mol2_filename, frcmod_filename])

    ffxml_filename = molecule_name + '.ffxml'
    create_ffxml_file([gaff_mol2_filename], [frcmod_filename], ffxml_filename)

    traj = md.load(gaff_mol2_filename)  # Read mol2 file.
    positions = traj.openmm_positions(0)  # Extract OpenMM-united positions of first (and only) trajectory frame
    topology = traj.top.to_openmm()

    # Create System object.
    forcefield = app.ForceField(ffxml_filename)
    system = forcefield.createSystem(topology, nonbondedMethod=app.NoCutoff, constraints=None, implicitSolvent=None)

    # Create integrator.
    timestep = 1.0 * units.femtoseconds
    integrator = simtk.openmm.VerletIntegrator(timestep)

    # Create simulation.
    platform = simtk.openmm.Platform.getPlatformByName("Reference")
    simulation = app.Simulation(topology, system, integrator, platform=platform)
    simulation.context.setPositions(positions)

    return simulation

def create_leap_simulation(molecule_name, gaff_mol2_filename, frcmod_filename):
    """Create an OpenMM simulation using a Gaff mol2 file and frcmod file.

    Parameters
    ----------
    molecule_name : str
        Name of the molecule
    gaff_mol2_filename : str
        Filename of input (GAFF!) mol2 file
    frcmod_filename : str
        Use this frcmod filename

    """

    # Parameterize system with LEaP.
    (prmtop_filename, inpcrd_filename) = run_tleap(molecule_name, gaff_mol2_filename, frcmod_filename)

    # Create System object.
    prmtop = app.AmberPrmtopFile(prmtop_filename)
    topology = prmtop.topology
    system = prmtop.createSystem(nonbondedMethod=app.NoCutoff, constraints=None, implicitSolvent=None)

    # Read positions.
    inpcrd = app.AmberInpcrdFile(inpcrd_filename)
    positions = inpcrd.getPositions()

    # Create integrator.
    timestep = 1.0 * units.femtoseconds
    integrator = simtk.openmm.VerletIntegrator(timestep)

    platform = simtk.openmm.Platform.getPlatformByName("Reference")
    simulation = app.Simulation(topology, system, integrator, platform=platform)
    simulation.context.setPositions(positions)

    return simulation

def test_molecule(molecule_name, tripos_mol2_filename, charge_method=None):
    """Create a GAFF molecule via LEAP and ffXML and compare force terms.

    Parameters
    ----------
    molecule_name : str
        Name of the molecule
    tripos_mol2_filename : str
        Filename of input mol2 file
    charge_method : str, default=None
        If None, use charges in existing MOL2.  Otherwise, use a charge
        model when running antechamber.
    """

    # Generate GAFF parameters.
    (gaff_mol2_filename, frcmod_filename) = run_antechamber(molecule_name, tripos_mol2_filename, charge_method=charge_method)

    # Create simulations.
    simulation_ffxml = create_ffxml_simulation(molecule_name, gaff_mol2_filename, frcmod_filename)
    simulation_leap  = create_leap_simulation(molecule_name, gaff_mol2_filename, frcmod_filename)

    # Compare simulations.
    syscheck = system_checker.SystemChecker(simulation_ffxml, simulation_leap)
    syscheck.check_force_parameters()

    groups0, groups1 = syscheck.check_energy_groups()
    energy0, energy1 = syscheck.check_energies()

def get_data_filename(relative_path):
    """Get the full path to one of the reference files shipped for testing

    In the source distribution, these files are in ``gaff2xml/chemicals/*/``,
    but on installation, they're moved to somewhere in the user's python
    site-packages directory.

    Parameters
    ----------
    name : str
        Name of the file to load (with respect to the gaff2xml folder).

    """

    fn = resource_filename('gaff2xml', relative_path)

    if not os.path.exists(fn):
        raise ValueError("Sorry! %s does not exist. If you just added it, you'll have to re-install" % fn)

    return fn

def smiles_to_mdtraj_ffxml(smiles_strings, base_molecule_name="lig"):
    """Generate an MDTraj object from a smiles string.

    Parameters
    ----------
    smiles_strings : list(str)
        Smiles strings to create molecules for
    base_molecule_name : str, optional, default='lig'
        Base name of molecule to use inside parameter files.

    Returns
    -------
    traj : mdtraj.Trajectory
        MDTraj object for molecule
    ffxml : StringIO
        StringIO representation of ffxml file.

    Notes
    -----
    ffxml can be directly input to OpenMM e.g.
    `forcefield = app.ForceField(ffxml)`
    """
    try:
        from rdkit import Chem
        from rdkit.Chem import AllChem
    except ImportError:
        raise(ImportError("Must install rdkit to use smiles conversion."))

    gaff_mol2_filenames = []
    frcmod_filenames = []
    trajectories = []
    for k, smiles_string in enumerate(smiles_strings):
        molecule_name = "%s-%d" % (base_molecule_name, k)
        m = Chem.MolFromSmiles(smiles_string)
        m = Chem.AddHs(m)
        AllChem.EmbedMolecule(m)
        AllChem.UFFOptimizeMolecule(m)

        pdb_filename = tempfile.mktemp(suffix=".pdb")
        Chem.MolToPDBFile(m, pdb_filename)

        mol2_filename = tempfile.mktemp(suffix=".mol2")

        convert_molecule(pdb_filename, mol2_filename)  # This is necessary because PDB double bonds are not handled by antechamber...
        print(mol2_filename)

        gaff_mol2_filename, frcmod_filename = run_antechamber(molecule_name, mol2_filename)
        traj = md.load(gaff_mol2_filename)
        print(gaff_mol2_filename)
        print(traj)

        for atom in traj.top.atoms:
            atom.residue.name = molecule_name

        gaff_mol2_filenames.append(gaff_mol2_filename)
        frcmod_filenames.append(frcmod_filename)
        trajectories.append(traj)

    ffxml = create_ffxml_file(gaff_mol2_filenames, frcmod_filenames, override_mol2_residue_name=molecule_name)

    return trajectories, ffxml

<|endoftext|>


<commit_before>import tornado
import tornado.gen

from req import Service
from req import ApiRequestHandler

class Testdata(ApiRequestHandler):
    @tornado.gen.coroutine
    def get(self, problem_id):
        err, res = yield from Service.Testdata.get_testdata_list({'id': problem_id})
        if err:
            self.render(err)
        else:
            self.render(res)

    @tornado.gen.coroutine
    def post(self, problem_id):
        args = ['score', 'time_limit', 'memory_limit', 'output_limit', 'input[file]', 'output[file]']
        data = self.get_args(args)
        data['problem_id'] = problem_id
        err, res = yield from Service.Testdata.post_testdata(data)
        if err:
            self.render(err)
        else:
            err, res = yield from Service.Testdata.get_testdata(res)
            self.render(res)

class Testdatum(ApiRequestHandler):
    @tornado.gen.coroutine
    def get(self, problem_id, testdata_id):
        err, res = yield from Service.Testdata.get_testdata({'id': testdata_id})
        if err:
            self.render(err)
        else:
            self.render(res)

    @tornado.gen.coroutine
    def put(self, problem_id, testdata_id):
        args = ['score', 'time_limit', 'memory_limit', 'output_limit', 'input[file]', 'output[file]']
        data = self.get_args(args)
        data['id'] = testdata_id
        err, res = yield from Service.Testdata.put_testdata(data)
        if err:
            self.render(err)
        else:
            self.render(res)

    @tornado.gen.coroutine
    def delete(self, problem_id, testdata_id):
        err, res = yield from Service.Testdata.delete_testdata({'id': testdata_id})
        if err:
            self.render(err)
        else:
            self.render(res)

<commit_msg>return testdata<commit_after>import tornado
import tornado.gen

from req import Service
from req import ApiRequestHandler

class Testdata(ApiRequestHandler):
    @tornado.gen.coroutine
    def get(self, problem_id):
        err, res = yield from Service.Testdata.get_testdata_list({'id': problem_id})
        if err:
            self.render(err)
        else:
            self.render(res)

    @tornado.gen.coroutine
    def post(self, problem_id):
        args = ['score', 'time_limit', 'memory_limit', 'output_limit', 'input[file]', 'output[file]']
        data = self.get_args(args)
        data['problem_id'] = problem_id
        err, res = yield from Service.Testdata.post_testdata(data)
        if err:
            self.render(err)
        else:
            err, res = yield from Service.Testdata.get_testdata(res)
            self.render(res)

class Testdatum(ApiRequestHandler):
    @tornado.gen.coroutine
    def get(self, problem_id, testdata_id):
        err, res = yield from Service.Testdata.get_testdata({'id': testdata_id})
        if err:
            self.render(err)
        else:
            self.render(res)

    @tornado.gen.coroutine
    def put(self, problem_id, testdata_id):
        args = ['score', 'time_limit', 'memory_limit', 'output_limit', 'input[file]', 'output[file]']
        data = self.get_args(args)
        data['id'] = testdata_id
        err, res = yield from Service.Testdata.put_testdata(data)
        if err:
            self.render(err)
        else:
            err, res = yield from Service.Testdata.get_testdata(data)
            self.render(res)

    @tornado.gen.coroutine
    def delete(self, problem_id, testdata_id):
        err, res = yield from Service.Testdata.delete_testdata({'id': testdata_id})
        if err:
            self.render(err)
        else:
            self.render(res)

<|endoftext|>


<commit_before>from django.conf import settings
from django.core.exceptions import ObjectDoesNotExist, ImproperlyConfigured
from django.core import validators
from django.db import backend, connection
from django.db.models.loading import get_apps, get_app, get_models, get_model, register_models
from django.db.models.query import Q
from django.db.models.manager import Manager
from django.db.models.base import Model, AdminOptions
from django.db.models.fields import *
from django.db.models.fields.related import ForeignKey, OneToOneField, ManyToManyField, ManyToOneRel, ManyToManyRel, OneToOneRel, TABULAR, STACKED
from django.db.models.fields.generic import GenericRelation, GenericRel, GenericForeignKey
from django.db.models import signals
from django.utils.functional import curry
from django.utils.text import capfirst

# Admin stages.
ADD, CHANGE, BOTH = 1, 2, 3

# Decorator. Takes a function that returns a tuple in this format:
#     (viewname, viewargs, viewkwargs)
# Returns a function that calls urlresolvers.reverse() on that data, to return
# the URL for those parameters.
def permalink(func):
    from django.core.urlresolvers import reverse
    def inner(*args, **kwargs):
        bits = func(*args, **kwargs)
        viewname = bits[0]
        return reverse(bits[0], None, *bits[1:3])
    return inner

class LazyDate(object):
    """
    Use in limit_choices_to to compare the field to dates calculated at run time
    instead of when the model is loaded.  For example::

        ... limit_choices_to = {'date__gt' : models.LazyDate(days=-3)} ...

    which will limit the choices to dates greater than three days ago.
    """
    def __init__(self, **kwargs):
        self.delta = datetime.timedelta(**kwargs)

    def __str__(self):
        return str(self.__get_value__())

    def __repr__(self):
        return "<LazyDate: %s>" % self.delta

    def __get_value__(self):
        return datetime.datetime.now() + self.delta

    def __getattr__(self, attr):
        return getattr(self.__get_value__(), attr)
<commit_msg>Fixed #2513 -- Changed LazyDate.__get_value__() to return date objects, not datetime objects<commit_after>from django.conf import settings
from django.core.exceptions import ObjectDoesNotExist, ImproperlyConfigured
from django.core import validators
from django.db import backend, connection
from django.db.models.loading import get_apps, get_app, get_models, get_model, register_models
from django.db.models.query import Q
from django.db.models.manager import Manager
from django.db.models.base import Model, AdminOptions
from django.db.models.fields import *
from django.db.models.fields.related import ForeignKey, OneToOneField, ManyToManyField, ManyToOneRel, ManyToManyRel, OneToOneRel, TABULAR, STACKED
from django.db.models.fields.generic import GenericRelation, GenericRel, GenericForeignKey
from django.db.models import signals
from django.utils.functional import curry
from django.utils.text import capfirst

# Admin stages.
ADD, CHANGE, BOTH = 1, 2, 3

# Decorator. Takes a function that returns a tuple in this format:
#     (viewname, viewargs, viewkwargs)
# Returns a function that calls urlresolvers.reverse() on that data, to return
# the URL for those parameters.
def permalink(func):
    from django.core.urlresolvers import reverse
    def inner(*args, **kwargs):
        bits = func(*args, **kwargs)
        viewname = bits[0]
        return reverse(bits[0], None, *bits[1:3])
    return inner

class LazyDate(object):
    """
    Use in limit_choices_to to compare the field to dates calculated at run time
    instead of when the model is loaded.  For example::

        ... limit_choices_to = {'date__gt' : models.LazyDate(days=-3)} ...

    which will limit the choices to dates greater than three days ago.
    """
    def __init__(self, **kwargs):
        self.delta = datetime.timedelta(**kwargs)

    def __str__(self):
        return str(self.__get_value__())

    def __repr__(self):
        return "<LazyDate: %s>" % self.delta

    def __get_value__(self):
        return (datetime.datetime.now() + self.delta).date()

    def __getattr__(self, attr):
        return getattr(self.__get_value__(), attr)
<|endoftext|>


<commit_before># SPDX-License-Identifier: MIT
# Copyright (C) 2019-2021 Tobias Gruetzmacher
# Copyright (C) 2019-2020 Daniel Ring
import json

from ..scraper import _ParserScraper

class MangaDex(_ParserScraper):
    multipleImagesPerStrip = True

    def __init__(self, name, mangaId):
        super(MangaDex, self).__init__('MangaDex/' + name)

        baseUrl = 'https://api.mangadex.org/'
        self.url = baseUrl + 'manga/%s' % mangaId
        self.chaptersUrl = baseUrl + 'manga/%s/feed?translatedLanguage[]=en&order[chapter]=desc&limit=500' % mangaId
        self.stripUrl = baseUrl + 'chapter/%s'
        self.cdnUrl = baseUrl + 'at-home/server/%s'
        self.imageUrl = 'https://s5.mangadex.org/data/%s/%%s'

    def starter(self):
        # Retrieve manga metadata from API
        mangaData = self.session.get(self.url)
        mangaData.raise_for_status()
        manga = mangaData.json()['data']

        # Retrieve chapter list from API
        chapterList = []
        chapterTotal = 1
        chapterOffset = 0
        while len(chapterList) < chapterTotal:
            chapterData = self.session.get(self.chaptersUrl + '&offset=%d' % chapterOffset)
            chapterData.raise_for_status()
            chapterBlock = chapterData.json()
            chapterTotal = chapterBlock['total']
            chapterOffset = chapterBlock['offset'] + chapterBlock['limit']
            chapterList.extend(chapterBlock['data'])

        # Determine if manga is complete and/or adult
        if manga['attributes']['lastChapter'] != '0':
            for chapter in chapterList:
                if chapter['attributes']['chapter'] == manga['attributes']['lastChapter']:
                    self.endOfLife = True
                    break

        if manga['attributes']['contentRating'] != 'safe':
            self.adult = True

        # Prepare chapter list
        self.chapters = []
        for chapter in chapterList:
            if chapter['attributes']['chapter'] == '':
                continue
            if chapter['attributes']['pages'] == 0:
                continue
            if len(self.chapters) >= 1 and chapter['attributes']['chapter'] == self.chapters[-1]['attributes']['chapter']:
                continue
            self.chapters.append(chapter)
        self.chapters.reverse()

        # Find first and last chapter
        self.firstStripUrl = self.stripUrl % self.chapters[0]['id']
        return self.stripUrl % self.chapters[-1]['id']

    def getPrevUrl(self, url, data):
        # Determine previous chapter ID from cached list
        chapterId = url.rsplit('/', 1)[-1]
        chapter = list(filter(lambda c: c['id'] == chapterId, self.chapters))
        if len(chapter) == 0:
            return None
        return self.stripUrl % self.chapters[self.chapters.index(chapter[0]) - 1]['id']

    def fetchUrls(self, url, data, urlSearch):
        # Retrieve chapter metadata from API
        chapterData = json.loads(data.text_content())
        self.chapter = chapterData['data']
        cdnData = self.session.get(self.cdnUrl % self.chapter['id'])
        cdnData.raise_for_status()
        cdnBlock = cdnData.json()

        # Save link order for position-based filenames
        imageUrl = self.imageUrl % cdnBlock['chapter']['hash']
        self.imageUrls = [imageUrl % page for page in cdnBlock['chapter']['data']]
        return self.imageUrls

    def namer(self, imageUrl, pageUrl):
        # Construct filename from episode number and page index in array
        chapterNum = self.chapter['attributes']['chapter'] if self.chapter['attributes']['chapter'] != None else 0
        pageNum = self.imageUrls.index(imageUrl)
        pageExt = imageUrl.rsplit('.')[-1]
        return '%s-%02d.%s' % (chapterNum, pageNum, pageExt)

    @classmethod
    def getmodules(cls):
        return (
            cls('AttackOnTitan', '304ceac3-8cdb-4fe7-acf7-2b6ff7a60613'),
            cls('Beastars', 'f5e3baad-3cd4-427c-a2ec-ad7d776b370d'),
            cls('BokuNoKokoroNoYabaiYatsu', '3df1a9a3-a1be-47a3-9e90-9b3e55b1d0ac'),
            cls('CheerfulAmnesia', 'f9448f90-c068-4b6a-8c85-03d739aef255'),
            cls('DoChokkyuuKareshiXKanojo', 'efb62763-c940-4495-aba5-69c192a999a4'),
            cls('DeliciousinDungeon', 'd90ea6cb-7bc3-4d80-8af0-28557e6c4e17'),
            cls('DragonDrive', '5c06ae70-b5cf-431a-bcd5-262a411de527'),
            cls('FuguushokuKajishiDakedoSaikyouDesu', '17b3b648-fd89-4a69-9a42-6068ffbfa7a7'),
            cls('GanbareDoukiChan', '190616bc-7da6-45fd-abd4-dd2ca656c183'),
            cls('HangingOutWithAGamerGirl', 'de9e3b62-eac5-4c0a-917d-ffccad694381'),
            cls('HoriMiya', 'a25e46ec-30f7-4db6-89df-cacbc1d9a900'),
            cls('HowToOpenATriangularRiceball', '6ebd90ce-d5e8-49c0-a4bc-e02e0f8ecb93'),
            cls('HunterXHunter', 'db692d58-4b13-4174-ae8c-30c515c0689c'),
            cls('IchaichasuruToOkaneGaWaichauFutariNoHanashi', '8eaaec7d-7aa7-490e-8d52-5a3d0a28e78b'),
            cls('InterspeciesReviewers', '1b2fddf9-1385-4f3c-b37a-cf86a9428b1a'),
            cls('JahySamaWaKujikenai', '2f4e5f5b-d930-4266-8c8a-c4cf9a81e51f'),
            cls('JingaiNoYomeToIchaIchaSuru', '809d2493-df3c-4e72-a57e-3e0026cae9fb'),
            cls('KaetteKudasaiAkutsuSan', '737a846b-2e67-4d63-9f7e-f54b3beebac4'),
            cls('KawaiiJoushiWoKomarasetai', '23b7cc7a-df89-4049-af28-1fa78f88713e'),
            cls('KanojoOkarishimasu', '32fdfe9b-6e11-4a13-9e36-dcd8ea77b4e4'),
            cls('KoiToUtatane', 'f7d40a27-e289-45b3-9c68-d1cb251897e6'),
            cls('Lv2KaraCheatDattaMotoYuushaKouhoNoMattariIsekaiLife', '58bc83a0-1808-484e-88b9-17e167469e23'),
            cls('MaouNoOreGaDoreiElfWoYomeNiShitandaGaDouMederebaIi', '55ace2fb-e157-4d76-9e72-67c6bd762a39'),
            cls('ModernMoGal', 'b1953f80-36f7-492c-b0f8-e9dd0ad01752'),
            cls('MyTinySenpaiFromWork', '28ed63af-61f8-43af-bac3-762030c72963'),
            cls('OMaidensinYourSavageSeason', 'c4613b7d-7a6e-48f9-82f0-bce3dd33383a'),
            cls('OokamiShounenWaKyouMoUsoOKasaneru', '5e77d9e2-2e44-431a-a995-5fefd411e55e'),
            cls('OokamiToKoshinryou', 'de900fd3-c94c-4148-bbcb-ca56eaeb57a4'),
            cls('OtomeYoukaiZakuro', 'c1fa97be-0f1f-4686-84bc-806881c97d53'),
            cls('OversimplifiedSCP', 'e911fe33-a9b3-43dc-9eb7-f5ee081a6dc8'),
            cls('PashiriNaBokuToKoisuruBanchouSan', '838e5b3a-51c8-44cf-b6e2-68193416f6fe'),
            cls('PleaseDontBullyMeNagatoro', 'd86cf65b-5f6c-437d-a0af-19a31f94ec55'),
            cls('PleaseDontBullyMeNagatoroComicAnthology', '2a4bc9ec-2d70-428a-8b46-27f6218ed267'),
            cls('PleaseTellMeGalkochan', '7a2f2f6b-a6a6-4149-879b-3fc2f6916549'),
            cls('RebuildWorld', '99182618-ae92-4aec-a5df-518659b7b613'),
            cls('SaekiSanWaNemutteru', 'd9aecdab-8aef-4b90-98d5-32e86faffb28'),
            cls('SeijoSamaIieToorisugariNoMamonotsukaiDesu', 'd4c40e73-251a-4bcb-a5a6-1edeec1e00e7'),
            cls('SenpaiGaUzaiKouhaiNoHanashi', 'af38f328-8df1-4b4c-a272-e737625c3ddc'),
            cls('SewayakiKitsuneNoSenkoSan', 'c26269c7-0f5d-4966-8cd5-b79acb86fb7a'),
            cls('SousouNoFrieren', 'b0b721ff-c388-4486-aa0f-c2b0bb321512'),
            cls('SwordArtOnline', '3dd0b814-23f4-4342-b75b-f206598534f6'),
            cls('SwordArtOnlineProgressive', '22ea3f54-11e4-4932-a527-89d63d3a62d9'),
            cls('TadokoroSan', '8ffbfa2f-23fa-4490-848e-942581a4d873'),
            cls('TamenDeGushi', '3f1453fb-9dac-4aca-a2ea-69613856c952'),
            cls('TenseiShitaraSlimeDattaKen', 'e78a489b-6632-4d61-b00b-5206f5b8b22b'),
            cls('TheNewGate', 'b41bef1e-7df9-4255-bd82-ecf570fec566'),
            cls('TheWolfAndRedRidingHood', 'a7d1283b-ed38-4659-b8bc-47bfca5ccb8a'),
            cls('TomoChanWaOnnanoko', '76ee7069-23b4-493c-bc44-34ccbf3051a8'),
            cls('TonikakuKawaii', '30f3ac69-21b6-45ad-a110-d011b7aaadaa'),
            cls('UramikoiKoiUramikoi', '009b6788-48f3-4e78-975c-097f54def7ab'),
            cls('YotsubaAnd', '58be6aa6-06cb-4ca5-bd20-f1392ce451fb'),
            cls('YuYuHakusho', '44a5cbe1-0204-4cc7-a1ff-0fda2ac004b6'),
        )
<commit_msg>Fix MangaDex site engine<commit_after># SPDX-License-Identifier: MIT
# Copyright (C) 2019-2021 Tobias Gruetzmacher
# Copyright (C) 2019-2020 Daniel Ring
import json

from ..scraper import _ParserScraper

class MangaDex(_ParserScraper):
    multipleImagesPerStrip = True

    def __init__(self, name, mangaId):
        super(MangaDex, self).__init__('MangaDex/' + name)

        baseUrl = 'https://api.mangadex.org/'
        self.url = baseUrl + 'manga/%s' % mangaId
        self.chaptersUrl = baseUrl + 'manga/%s/feed?translatedLanguage[]=en&order[chapter]=desc&limit=500' % mangaId
        self.stripUrl = baseUrl + 'chapter/%s'
        self.cdnUrl = baseUrl + 'at-home/server/%s'
        self.imageUrl = 'https://uploads.mangadex.org/data/%s/%%s'

    def starter(self):
        # Retrieve manga metadata from API
        mangaData = self.session.get(self.url)
        mangaData.raise_for_status()
        manga = mangaData.json()['data']

        # Retrieve chapter list from API
        chapterList = []
        chapterTotal = 1
        chapterOffset = 0
        while len(chapterList) < chapterTotal:
            chapterData = self.session.get(self.chaptersUrl + '&offset=%d' % chapterOffset)
            chapterData.raise_for_status()
            chapterBlock = chapterData.json()
            chapterTotal = chapterBlock['total']
            chapterOffset = chapterBlock['offset'] + chapterBlock['limit']
            chapterList.extend(chapterBlock['data'])

        # Determine if manga is complete and/or adult
        if manga['attributes']['lastChapter'] != '0':
            for chapter in chapterList:
                if chapter['attributes']['chapter'] == manga['attributes']['lastChapter']:
                    self.endOfLife = True
                    break

        if manga['attributes']['contentRating'] != 'safe':
            self.adult = True

        # Prepare chapter list
        self.chapters = []
        for chapter in chapterList:
            if chapter['attributes']['chapter'] == '':
                continue
            if chapter['attributes']['pages'] == 0:
                continue
            if len(self.chapters) >= 1 and chapter['attributes']['chapter'] == self.chapters[-1]['attributes']['chapter']:
                continue
            self.chapters.append(chapter)
        self.chapters.reverse()

        # Find first and last chapter
        self.firstStripUrl = self.stripUrl % self.chapters[0]['id']
        return self.stripUrl % self.chapters[-1]['id']

    def getPrevUrl(self, url, data):
        # Determine previous chapter ID from cached list
        chapterId = url.rsplit('/', 1)[-1]
        chapter = list(filter(lambda c: c['id'] == chapterId, self.chapters))
        if len(chapter) == 0:
            return None
        return self.stripUrl % self.chapters[self.chapters.index(chapter[0]) - 1]['id']

    def fetchUrls(self, url, data, urlSearch):
        # Retrieve chapter metadata from API
        chapterData = json.loads(data.text_content())
        self.chapter = chapterData['data']
        cdnData = self.session.get(self.cdnUrl % self.chapter['id'])
        cdnData.raise_for_status()
        cdnBlock = cdnData.json()

        # Save link order for position-based filenames
        imageUrl = self.imageUrl % cdnBlock['chapter']['hash']
        self.imageUrls = [imageUrl % page for page in cdnBlock['chapter']['data']]
        return self.imageUrls

    def namer(self, imageUrl, pageUrl):
        # Construct filename from episode number and page index in array
        chapterNum = self.chapter['attributes']['chapter'] if self.chapter['attributes']['chapter'] != None else 0
        pageNum = self.imageUrls.index(imageUrl)
        pageExt = imageUrl.rsplit('.')[-1]
        return '%s-%02d.%s' % (chapterNum, pageNum, pageExt)

    @classmethod
    def getmodules(cls):
        return (
            cls('AttackOnTitan', '304ceac3-8cdb-4fe7-acf7-2b6ff7a60613'),
            cls('Beastars', 'f5e3baad-3cd4-427c-a2ec-ad7d776b370d'),
            cls('BokuNoKokoroNoYabaiYatsu', '3df1a9a3-a1be-47a3-9e90-9b3e55b1d0ac'),
            cls('CheerfulAmnesia', 'f9448f90-c068-4b6a-8c85-03d739aef255'),
            cls('DoChokkyuuKareshiXKanojo', 'efb62763-c940-4495-aba5-69c192a999a4'),
            cls('DeliciousinDungeon', 'd90ea6cb-7bc3-4d80-8af0-28557e6c4e17'),
            cls('DragonDrive', '5c06ae70-b5cf-431a-bcd5-262a411de527'),
            cls('FuguushokuKajishiDakedoSaikyouDesu', '17b3b648-fd89-4a69-9a42-6068ffbfa7a7'),
            cls('GanbareDoukiChan', '190616bc-7da6-45fd-abd4-dd2ca656c183'),
            cls('HangingOutWithAGamerGirl', 'de9e3b62-eac5-4c0a-917d-ffccad694381'),
            cls('HoriMiya', 'a25e46ec-30f7-4db6-89df-cacbc1d9a900'),
            cls('HowToOpenATriangularRiceball', '6ebd90ce-d5e8-49c0-a4bc-e02e0f8ecb93'),
            cls('HunterXHunter', 'db692d58-4b13-4174-ae8c-30c515c0689c'),
            cls('IchaichasuruToOkaneGaWaichauFutariNoHanashi', '8eaaec7d-7aa7-490e-8d52-5a3d0a28e78b'),
            cls('InterspeciesReviewers', '1b2fddf9-1385-4f3c-b37a-cf86a9428b1a'),
            cls('JahySamaWaKujikenai', '2f4e5f5b-d930-4266-8c8a-c4cf9a81e51f'),
            cls('JingaiNoYomeToIchaIchaSuru', '809d2493-df3c-4e72-a57e-3e0026cae9fb'),
            cls('KaetteKudasaiAkutsuSan', '737a846b-2e67-4d63-9f7e-f54b3beebac4'),
            cls('KawaiiJoushiWoKomarasetai', '23b7cc7a-df89-4049-af28-1fa78f88713e'),
            cls('KanojoOkarishimasu', '32fdfe9b-6e11-4a13-9e36-dcd8ea77b4e4'),
            cls('KoiToUtatane', 'f7d40a27-e289-45b3-9c68-d1cb251897e6'),
            cls('Lv2KaraCheatDattaMotoYuushaKouhoNoMattariIsekaiLife', '58bc83a0-1808-484e-88b9-17e167469e23'),
            cls('MaouNoOreGaDoreiElfWoYomeNiShitandaGaDouMederebaIi', '55ace2fb-e157-4d76-9e72-67c6bd762a39'),
            cls('ModernMoGal', 'b1953f80-36f7-492c-b0f8-e9dd0ad01752'),
            cls('MyTinySenpaiFromWork', '28ed63af-61f8-43af-bac3-762030c72963'),
            cls('OMaidensinYourSavageSeason', 'c4613b7d-7a6e-48f9-82f0-bce3dd33383a'),
            cls('OokamiShounenWaKyouMoUsoOKasaneru', '5e77d9e2-2e44-431a-a995-5fefd411e55e'),
            cls('OokamiToKoshinryou', 'de900fd3-c94c-4148-bbcb-ca56eaeb57a4'),
            cls('OtomeYoukaiZakuro', 'c1fa97be-0f1f-4686-84bc-806881c97d53'),
            cls('OversimplifiedSCP', 'e911fe33-a9b3-43dc-9eb7-f5ee081a6dc8'),
            cls('PashiriNaBokuToKoisuruBanchouSan', '838e5b3a-51c8-44cf-b6e2-68193416f6fe'),
            cls('PleaseDontBullyMeNagatoro', 'd86cf65b-5f6c-437d-a0af-19a31f94ec55'),
            cls('PleaseDontBullyMeNagatoroComicAnthology', '2a4bc9ec-2d70-428a-8b46-27f6218ed267'),
            cls('PleaseTellMeGalkochan', '7a2f2f6b-a6a6-4149-879b-3fc2f6916549'),
            cls('RebuildWorld', '99182618-ae92-4aec-a5df-518659b7b613'),
            cls('SaekiSanWaNemutteru', 'd9aecdab-8aef-4b90-98d5-32e86faffb28'),
            cls('SeijoSamaIieToorisugariNoMamonotsukaiDesu', 'd4c40e73-251a-4bcb-a5a6-1edeec1e00e7'),
            cls('SenpaiGaUzaiKouhaiNoHanashi', 'af38f328-8df1-4b4c-a272-e737625c3ddc'),
            cls('SewayakiKitsuneNoSenkoSan', 'c26269c7-0f5d-4966-8cd5-b79acb86fb7a'),
            cls('SousouNoFrieren', 'b0b721ff-c388-4486-aa0f-c2b0bb321512'),
            cls('SwordArtOnline', '3dd0b814-23f4-4342-b75b-f206598534f6'),
            cls('SwordArtOnlineProgressive', '22ea3f54-11e4-4932-a527-89d63d3a62d9'),
            cls('TadokoroSan', '8ffbfa2f-23fa-4490-848e-942581a4d873'),
            cls('TamenDeGushi', '3f1453fb-9dac-4aca-a2ea-69613856c952'),
            cls('TenseiShitaraSlimeDattaKen', 'e78a489b-6632-4d61-b00b-5206f5b8b22b'),
            cls('TheNewGate', 'b41bef1e-7df9-4255-bd82-ecf570fec566'),
            cls('TheWolfAndRedRidingHood', 'a7d1283b-ed38-4659-b8bc-47bfca5ccb8a'),
            cls('TomoChanWaOnnanoko', '76ee7069-23b4-493c-bc44-34ccbf3051a8'),
            cls('TonikakuKawaii', '30f3ac69-21b6-45ad-a110-d011b7aaadaa'),
            cls('UramikoiKoiUramikoi', '009b6788-48f3-4e78-975c-097f54def7ab'),
            cls('YotsubaAnd', '58be6aa6-06cb-4ca5-bd20-f1392ce451fb'),
            cls('YuYuHakusho', '44a5cbe1-0204-4cc7-a1ff-0fda2ac004b6'),
        )
<|endoftext|>


<commit_before>import os
import sys

from collections import defaultdict, OrderedDict

from marisa_trie import BytesTrie

this_dir = os.path.realpath(os.path.dirname(__file__))
sys.path.append(os.path.realpath(os.path.join(os.pardir, os.pardir)))

sys.path.append(os.path.realpath(os.path.join(os.pardir, os.pardir, os.pardir, 'python')))

from geodata.encoding import safe_decode
from geodata.i18n.unicode_paths import DATA_DIR
from geodata.i18n.normalize import strip_accents
from geodata.i18n.unicode_properties import get_chars_by_script, get_script_languages
from address_normalizer.text.normalize import PhraseFilter
from address_normalizer.text.tokenize import *

WELL_REPRESENTED_LANGUAGES = set(['en', 'fr', 'it', 'de', 'nl', 'es', 'pt'])

# For toponyms, we want to limit the countries we consider to those where
# we the place names can themselves be considered training examples of the language
WELL_REPRESENTED_LANGUAGE_COUNTRIES = {
    'en': set(['gb', 'us', 'ca', 'au', 'nz', 'ie']),
    'fr': set(['fr']),
    'it': set(['it']),
    'de': set(['de', 'at']),
    'nl': set(['nl']),
    'es': set(['es', 'ar', 'mx', 'cl', 'co', 'pe', 'ec', 'pr', 'uy',
               've', 'cu', 'do', 'bo', 'gt', 'cr', 'py', 'sv', 'pa',
               'ni', 'hn']),
    'pt': set(['pt', 'br']),
}

DICTIONARIES_DIR = os.path.join(DATA_DIR, 'dictionaries')

PREFIX_KEY = u'\x02'
SUFFIX_KEY = u'\x03'

POSSIBLE_ROMAN_NUMERALS = set(['i', 'ii', 'iii', 'iv', 'v', 'vi', 'vii', 'viii', 'ix',
                               'x', 'xi', 'xii', 'xiii', 'xiv', 'xv', 'xvi', 'xvii', 'xviii', 'xix',
                               'xx', 'xxx', 'xl', 'l', 'lx', 'lxx', 'lxxx', 'xc',
                               'c', 'cc', 'ccc', 'cd', 'd', 'dc', 'dcc', 'dccc', 'cm',
                               'm', 'mm', 'mmm', 'mmmm'])

class DictionaryPhraseFilter(PhraseFilter):
    def __init__(self, *dictionaries):
        self.dictionaries = dictionaries

    def serialize(self, s):
        return s

    def deserialize(self, s):
        return s

    def configure(self, base_dir=DICTIONARIES_DIR):
        kvs = defaultdict(OrderedDict)
        for lang in os.listdir(DICTIONARIES_DIR):
            for filename in self.dictionaries:
                is_suffix_dictionary = 'suffixes' in filename
                is_prefix_dictionary = 'prefixes' in filename
                is_street_types_dictionary = 'street_types' in filename
                is_stopword_dictionary = 'stopwords' in filename

                path = os.path.join(DICTIONARIES_DIR, lang, filename)
                if not os.path.exists(path):
                    continue

                for line in open(path):
                    line = line.strip()
                    if not line:
                        continue

                    phrases = safe_decode(line).split(u'|')
                    if not phrases:
                        continue
                    canonical = strip_accents(phrases[0])

                    for phrase in phrases:

                        if phrase in POSSIBLE_ROMAN_NUMERALS:
                            continue

                        is_canonical = strip_accents(phrase) == canonical

                        if is_suffix_dictionary:
                            phrase = SUFFIX_KEY + phrase[::-1]
                        elif is_prefix_dictionary:
                            phrase = PREFIX_KEY + phrase

                        if is_canonical or is_street_types_dictionary or is_prefix_dictionary or is_suffix_dictionary:
                            kvs[phrase][lang] = (is_canonical, is_stopword_dictionary)

        kvs = [(k, '|'.join([v, str(int(c)), str(int(s))])) for k, vals in kvs.iteritems() for v, (c, s) in vals.iteritems()]

        self.trie = BytesTrie(kvs)
        self.configured = True

    def search_substring(self, s):
        if len(s) == 0:
            return None, 0

        for i in xrange(len(s) + 1):
            if not self.trie.has_keys_with_prefix(s[:i]):
                i -= 1
                break
        if i > 0:
            return (self.trie.get(s[:i]), i)
        else:
            return None, 0

    def search_suffix(self, token):
        suffix_search, suffix_len = self.search_substring(SUFFIX_KEY + token[::-1])
        if suffix_len > 0:
            suffix_len -= len(SUFFIX_KEY)
        return suffix_search, suffix_len

    def search_prefix(self, token):
        prefix_search, prefix_len = self.search_substring(PREFIX_KEY + token)
        if prefix_len > 0:
            prefix_len -= len(PREFIX_KEY)
        return prefix_search, prefix_len

    def basic_filter(self, tokens):
        return super(DictionaryPhraseFilter, self).filter(tokens)

    def filter(self, tokens):
        for c, t, data in self.basic_filter(tokens):
            if c != token_types.PHRASE:
                token = t[1]
                token_len = len(token)

                suffix_search, suffix_len = self.search_suffix(token)
                if suffix_search and self.trie.get(token[(token_len - suffix_len):].rstrip('.')):
                    yield (token_types.PHRASE, [(c,) + t], suffix_search)
                    continue
                prefix_search, prefix_len = self.search_prefix(token)
                if prefix_search and self.trie.get(token[:prefix_len]):
                    yield (token_types.PHRASE, [(c,) + t], prefix_search)
                    continue
            yield c, t, data

street_types_gazetteer = DictionaryPhraseFilter('street_types.txt',
                                                'directionals.txt',
                                                'concatenated_suffixes_separable.txt',
                                                'concatenated_suffixes_inseparable.txt',
                                                'concatenated_prefixes_separable.txt',
                                                'stopwords.txt',)

char_scripts = get_chars_by_script()
script_languages = {script: set(langs) for script, langs in get_script_languages().iteritems()}

UNKNOWN_SCRIPT = 'Unknown'
COMMON_SCRIPT = 'Common'
MAX_ASCII = 127

def get_string_script(s):
    s = safe_decode(s)
    script = last_script = UNKNOWN_SCRIPT
    is_ascii = True
    script_len = 0
    for c in s:
        if (ord(c)) < len(char_scripts):
            script = char_scripts[ord(c)]
        else:
            script = UNKNOWN_SCRIPT
        if script == COMMON_SCRIPT and last_script != UNKNOWN_SCRIPT:
            script = last_script
        if last_script != script and last_script != UNKNOWN_SCRIPT and last_script != COMMON_SCRIPT:
            break
        is_ascii = is_ascii and ord(c) <= MAX_ASCII
        script_len += 1
        if script != UNKNOWN_SCRIPT:
            last_script = script
    return (last_script, script_len, is_ascii)

LATIN_SCRIPT = 'Latin'
UNKNOWN_LANGUAGE = 'unk'
AMBIGUOUS_LANGUAGE = 'xxx'

def disambiguate_language(text, languages):
    valid_languages = OrderedDict(languages)
    script_langs = {}
    read_len = 0
    while read_len < len(text):
        script, script_len, is_ascii = get_string_script(text[read_len:])
        if script != LATIN_SCRIPT:
            script_langs[script] = set([l for l, d in languages if l in script_languages.get(script, [])])
        read_len += script_len

    num_defaults = sum((1 for lang, default in valid_languages.iteritems() if default))
    tokens = [(c, t.rstrip('.')) for c, t in tokenize(safe_decode(text).replace(u'-', u' ').lower())]

    current_lang = None
    possible_lang = None

    seen_languages = set()

    for c, t, data in street_types_gazetteer.filter(tokens):
        if c == token_types.PHRASE:
            valid = []
            data = [d.split('|') for d in data]
            potentials = [l for l, c, s in data if l in valid_languages]

            for lang, canonical, stopword in data:
                canonical = int(canonical)
                stopword = int(stopword)
                if lang not in valid_languages or (stopword and len(potentials) > 1):
                    continue
                is_default = valid_languages[lang]

                lang_valid = is_default or not seen_languages or lang in seen_languages

                if lang_valid and ((canonical and not stopword) or (is_default and len(potentials) == 1)):
                    valid.append(lang)
                elif is_default and num_defaults > 1 and current_lang is not None and current_lang != lang:
                    return AMBIGUOUS_LANGUAGE
                elif stopword and canonical and not is_default and lang in seen_languages:
                    valid.append(lang)
                elif not seen_languages and len(potentials) == 1 and len(t[0][1]) > 1:
                    possible_lang = lang if possible_lang is None or possible_lang == lang else None

            if seen_languages and valid and not any((l in seen_languages for l in valid)) and \
               (not any((valid_languages.get(l) for l in valid)) or any((valid_languages.get(l) for l in seen_languages))):
                return AMBIGUOUS_LANGUAGE

            if len(valid) == 1:
                current_lang = valid[0]
            else:
                valid_default = [l for l in valid if valid_languages.get(l)]
                if len(valid_default) == 1 and current_lang is not None and valid_default[0] != current_lang:
                    return AMBIGUOUS_LANGUAGE
                elif len(valid_default) == 1:
                    current_lang = valid_default[0]

            if any((current_lang not in langs for script, langs in script_langs.iteritems())):
                return AMBIGUOUS_LANGUAGE

            seen_languages.update(valid)

    if current_lang is not None:
        return current_lang
    elif possible_lang is not None:
        if not any((possible_lang not in langs for script, langs in script_langs.iteritems())):
            return possible_lang
        else:
            return AMBIGUOUS_LANGUAGE
    return UNKNOWN_LANGUAGE
<commit_msg>[unicode] Adding wide char support for language disambiguation (comes up in venue names), despite the likelihood of running on a narrow Python build. Rolling back common script chars at a script break, so in the case of e.g. Cyrllic name (Latin name), the segmentation is done at the space before the paren.<commit_after>import os
import sys

from collections import defaultdict, OrderedDict

from marisa_trie import BytesTrie

this_dir = os.path.realpath(os.path.dirname(__file__))
sys.path.append(os.path.realpath(os.path.join(os.pardir, os.pardir)))

sys.path.append(os.path.realpath(os.path.join(os.pardir, os.pardir, os.pardir, 'python')))

from geodata.encoding import safe_decode
from geodata.i18n.unicode_paths import DATA_DIR
from geodata.i18n.normalize import strip_accents
from geodata.i18n.unicode_properties import get_chars_by_script, get_script_languages
from address_normalizer.text.normalize import PhraseFilter
from address_normalizer.text.tokenize import *

WELL_REPRESENTED_LANGUAGES = set(['en', 'fr', 'it', 'de', 'nl', 'es', 'pt'])

# For toponyms, we want to limit the countries we consider to those where
# we the place names can themselves be considered training examples of the language
WELL_REPRESENTED_LANGUAGE_COUNTRIES = {
    'en': set(['gb', 'us', 'ca', 'au', 'nz', 'ie']),
    'fr': set(['fr']),
    'it': set(['it']),
    'de': set(['de', 'at']),
    'nl': set(['nl']),
    'es': set(['es', 'ar', 'mx', 'cl', 'co', 'pe', 'ec', 'pr', 'uy',
               've', 'cu', 'do', 'bo', 'gt', 'cr', 'py', 'sv', 'pa',
               'ni', 'hn']),
    'pt': set(['pt', 'br']),
}

DICTIONARIES_DIR = os.path.join(DATA_DIR, 'dictionaries')

PREFIX_KEY = u'\x02'
SUFFIX_KEY = u'\x03'

POSSIBLE_ROMAN_NUMERALS = set(['i', 'ii', 'iii', 'iv', 'v', 'vi', 'vii', 'viii', 'ix',
                               'x', 'xi', 'xii', 'xiii', 'xiv', 'xv', 'xvi', 'xvii', 'xviii', 'xix',
                               'xx', 'xxx', 'xl', 'l', 'lx', 'lxx', 'lxxx', 'xc',
                               'c', 'cc', 'ccc', 'cd', 'd', 'dc', 'dcc', 'dccc', 'cm',
                               'm', 'mm', 'mmm', 'mmmm'])

class DictionaryPhraseFilter(PhraseFilter):
    def __init__(self, *dictionaries):
        self.dictionaries = dictionaries

    def serialize(self, s):
        return s

    def deserialize(self, s):
        return s

    def configure(self, base_dir=DICTIONARIES_DIR):
        kvs = defaultdict(OrderedDict)
        for lang in os.listdir(DICTIONARIES_DIR):
            for filename in self.dictionaries:
                is_suffix_dictionary = 'suffixes' in filename
                is_prefix_dictionary = 'prefixes' in filename
                is_street_types_dictionary = 'street_types' in filename
                is_stopword_dictionary = 'stopwords' in filename

                path = os.path.join(DICTIONARIES_DIR, lang, filename)
                if not os.path.exists(path):
                    continue

                for line in open(path):
                    line = line.strip()
                    if not line:
                        continue

                    phrases = safe_decode(line).split(u'|')
                    if not phrases:
                        continue
                    canonical = strip_accents(phrases[0])

                    for phrase in phrases:

                        if phrase in POSSIBLE_ROMAN_NUMERALS:
                            continue

                        is_canonical = strip_accents(phrase) == canonical

                        if is_suffix_dictionary:
                            phrase = SUFFIX_KEY + phrase[::-1]
                        elif is_prefix_dictionary:
                            phrase = PREFIX_KEY + phrase

                        if is_canonical or is_street_types_dictionary or is_prefix_dictionary or is_suffix_dictionary:
                            kvs[phrase][lang] = (is_canonical, is_stopword_dictionary)

        kvs = [(k, '|'.join([v, str(int(c)), str(int(s))])) for k, vals in kvs.iteritems() for v, (c, s) in vals.iteritems()]

        self.trie = BytesTrie(kvs)
        self.configured = True

    def search_substring(self, s):
        if len(s) == 0:
            return None, 0

        for i in xrange(len(s) + 1):
            if not self.trie.has_keys_with_prefix(s[:i]):
                i -= 1
                break
        if i > 0:
            return (self.trie.get(s[:i]), i)
        else:
            return None, 0

    def search_suffix(self, token):
        suffix_search, suffix_len = self.search_substring(SUFFIX_KEY + token[::-1])
        if suffix_len > 0:
            suffix_len -= len(SUFFIX_KEY)
        return suffix_search, suffix_len

    def search_prefix(self, token):
        prefix_search, prefix_len = self.search_substring(PREFIX_KEY + token)
        if prefix_len > 0:
            prefix_len -= len(PREFIX_KEY)
        return prefix_search, prefix_len

    def basic_filter(self, tokens):
        return super(DictionaryPhraseFilter, self).filter(tokens)

    def filter(self, tokens):
        for c, t, data in self.basic_filter(tokens):
            if c != token_types.PHRASE:
                token = t[1]
                token_len = len(token)

                suffix_search, suffix_len = self.search_suffix(token)
                if suffix_search and self.trie.get(token[(token_len - suffix_len):].rstrip('.')):
                    yield (token_types.PHRASE, [(c,) + t], suffix_search)
                    continue
                prefix_search, prefix_len = self.search_prefix(token)
                if prefix_search and self.trie.get(token[:prefix_len]):
                    yield (token_types.PHRASE, [(c,) + t], prefix_search)
                    continue
            yield c, t, data

street_types_gazetteer = DictionaryPhraseFilter('street_types.txt',
                                                'directionals.txt',
                                                'concatenated_suffixes_separable.txt',
                                                'concatenated_suffixes_inseparable.txt',
                                                'concatenated_prefixes_separable.txt',
                                                'stopwords.txt',)

char_scripts = get_chars_by_script()
script_languages = {script: set(langs) for script, langs in get_script_languages().iteritems()}

UNKNOWN_SCRIPT = 'Unknown'
COMMON_SCRIPT = 'Common'
MAX_ASCII = 127

def get_string_script(s):
    s = safe_decode(s)
    str_len = len(s)
    script = last_script = UNKNOWN_SCRIPT
    is_ascii = True
    script_len = 0
    for c in wide_iter(s):
        script = char_scripts[wide_ord(c)]

        if script == COMMON_SCRIPT and last_script != UNKNOWN_SCRIPT:
            script = last_script
        if last_script != script and last_script != UNKNOWN_SCRIPT and last_script != COMMON_SCRIPT:
            if (script_len < str_len):
                for c in reversed(list(wide_iter(s[:script_len]))):
                    if char_scripts[wide_ord(c)] == COMMON_SCRIPT:
                        script_len -= 1
            break
        is_ascii = is_ascii and ord(c) <= MAX_ASCII
        script_len += 1
        if script != UNKNOWN_SCRIPT:
            last_script = script
    return (last_script, script_len, is_ascii)

LATIN_SCRIPT = 'Latin'
UNKNOWN_LANGUAGE = 'unk'
AMBIGUOUS_LANGUAGE = 'xxx'

def disambiguate_language(text, languages):
    valid_languages = OrderedDict(languages)
    script_langs = {}
    read_len = 0
    while read_len < len(text):
        script, script_len, is_ascii = get_string_script(text[read_len:])
        if script != LATIN_SCRIPT:
            script_langs[script] = set([l for l, d in languages if l in script_languages.get(script, [])])
        read_len += script_len

    num_defaults = sum((1 for lang, default in valid_languages.iteritems() if default))
    tokens = [(c, t.rstrip('.')) for c, t in tokenize(safe_decode(text).replace(u'-', u' ').lower())]

    current_lang = None
    possible_lang = None

    seen_languages = set()

    for c, t, data in street_types_gazetteer.filter(tokens):
        if c == token_types.PHRASE:
            valid = []
            data = [d.split('|') for d in data]
            potentials = [l for l, c, s in data if l in valid_languages]

            for lang, canonical, stopword in data:
                canonical = int(canonical)
                stopword = int(stopword)
                if lang not in valid_languages or (stopword and len(potentials) > 1):
                    continue
                is_default = valid_languages[lang]

                lang_valid = is_default or not seen_languages or lang in seen_languages

                if lang_valid and ((canonical and not stopword) or (is_default and len(potentials) == 1)):
                    valid.append(lang)
                elif is_default and num_defaults > 1 and current_lang is not None and current_lang != lang:
                    return AMBIGUOUS_LANGUAGE
                elif stopword and canonical and not is_default and lang in seen_languages:
                    valid.append(lang)
                elif not seen_languages and len(potentials) == 1 and len(t[0][1]) > 1:
                    possible_lang = lang if possible_lang is None or possible_lang == lang else None

            if seen_languages and valid and not any((l in seen_languages for l in valid)) and \
               (not any((valid_languages.get(l) for l in valid)) or any((valid_languages.get(l) for l in seen_languages))):
                return AMBIGUOUS_LANGUAGE

            if len(valid) == 1:
                current_lang = valid[0]
            else:
                valid_default = [l for l in valid if valid_languages.get(l)]
                if len(valid_default) == 1 and current_lang is not None and valid_default[0] != current_lang:
                    return AMBIGUOUS_LANGUAGE
                elif len(valid_default) == 1:
                    current_lang = valid_default[0]

            if any((current_lang not in langs for script, langs in script_langs.iteritems())):
                return AMBIGUOUS_LANGUAGE

            seen_languages.update(valid)

    if current_lang is not None:
        return current_lang
    elif possible_lang is not None:
        if not any((possible_lang not in langs for script, langs in script_langs.iteritems())):
            return possible_lang
        else:
            return AMBIGUOUS_LANGUAGE
    return UNKNOWN_LANGUAGE
<|endoftext|>


<commit_before># -*- coding: utf-8; -*-
#
# This file is part of Superdesk.
#
# Copyright 2013, 2014, 2015 Sourcefabric z.u. and contributors.
#
# For the full copyright and license information, please see the
# AUTHORS and LICENSE files distributed with this source code, or
# at https://www.sourcefabric.org/superdesk/license

from superdesk import get_resource_service, editor_utils
from superdesk.media.crop import CropService
from superdesk.metadata.item import ITEM_STATE, EMBARGO, SCHEDULE_SETTINGS
from superdesk.utc import utcnow
from superdesk.text_utils import update_word_count
from apps.archive.common import set_sign_off, ITEM_OPERATION, get_user
from apps.archive.archive import flush_renditions
from apps.tasks import send_to
from apps.auth import get_user_id
from .common import BasePublishService, BasePublishResource, ITEM_CORRECT
from superdesk.emails import send_translation_changed
from superdesk.activity import add_activity
from flask import g, current_app as app

def send_translation_notifications(original):
    translated_items = get_resource_service("published").find(
        {"translation_id": original.get("guid"), "last_published_version": True}
    )

    user_ids = set()
    for translated_item in translated_items:
        if translated_item.get("item_id") == original.get("_id"):
            changed_article = translated_item
            continue
        user_ids.add(translated_item.get("original_creator"))
        user_ids.add(translated_item.get("version_creator"))

    if len(user_ids) == 0 or changed_article is None:
        return

    add_activity("translated:changed", "", resource=None, item=changed_article, notify=user_ids)

    recipients = []
    for user_id in user_ids:
        user = get_resource_service("users").find_one(req=None, _id=user_id)
        if user is not None and user.get("email", None) is not None:
            recipients.append(user.get("email"))

    if len(recipients) == 0:
        return

    username = g.user.get("display_name") or g.user.get("username")
    send_translation_changed(username, changed_article, recipients)

class CorrectPublishResource(BasePublishResource):
    def __init__(self, endpoint_name, app, service):
        super().__init__(endpoint_name, app, service, "correct")

class CorrectPublishService(BasePublishService):
    publish_type = "correct"
    published_state = "corrected"
    item_operation = ITEM_CORRECT

    def set_state(self, original, updates):
        updates[ITEM_STATE] = self.published_state

        if original.get(EMBARGO) or updates.get(EMBARGO):
            # embargo time elapsed
            utc_embargo = updates.get(SCHEDULE_SETTINGS, {}).get("utc_{}".format(EMBARGO)) or original.get(
                SCHEDULE_SETTINGS, {}
            ).get("utc_{}".format(EMBARGO))
            if utc_embargo and utc_embargo < utcnow():
                # remove embargo information. so the next correction is without embargo.
                updates[EMBARGO] = None
                super().set_state(original, updates)
        else:
            super().set_state(original, updates)

    def change_being_corrected_to_published(self, updates, original):
        if app.config.get("CORRECTIONS_WORKFLOW") and original.get("state") == "correction":
            publish_service = get_resource_service("published")
            being_corrected_article = publish_service.find_one(
                req=None, guid=original.get("guid"), state="being_corrected"
            )

            if being_corrected_article.get("correction_sequence", 0) > 0:
                publish_service.patch(being_corrected_article["_id"], updates={"state": "corrected"})
            else:
                publish_service.patch(being_corrected_article["_id"], updates={"state": "published"})

    def send_to_original_desk(self, updates, original):
        if (
            app.config.get("CORRECTIONS_WORKFLOW")
            and original.get("state") == "correction"
            and original.get("task", {}).get("desk_history")
        ):
            send_to(
                doc=updates,
                desk_id=(original["task"]["desk_history"][0]),
                default_stage="working_stage",
                user_id=get_user_id(),
            )

    def on_update(self, updates, original):
        CropService().validate_multiple_crops(updates, original)
        super().on_update(updates, original)
        updates[ITEM_OPERATION] = self.item_operation
        updates["versioncreated"] = utcnow()
        updates["correction_sequence"] = original.get("correction_sequence", 1) + 1
        set_sign_off(updates, original)
        update_word_count(updates, original)
        flush_renditions(updates, original)
        self.change_being_corrected_to_published(updates, original)
        self.send_to_original_desk(updates, original)

    def update(self, id, updates, original):
        editor_utils.generate_fields(updates)
        get_resource_service("archive")._handle_media_updates(updates, original, get_user())
        super().update(id, updates, original)

    def on_updated(self, updates, original):
        super().on_updated(updates, original)
        send_translation_notifications(original)
<commit_msg>change the corrections workflow to show items with corrected state in the Desk output where they were published [SDCP-516] (#2197)<commit_after># -*- coding: utf-8; -*-
#
# This file is part of Superdesk.
#
# Copyright 2013, 2014, 2015 Sourcefabric z.u. and contributors.
#
# For the full copyright and license information, please see the
# AUTHORS and LICENSE files distributed with this source code, or
# at https://www.sourcefabric.org/superdesk/license

from superdesk import get_resource_service, editor_utils
from superdesk.media.crop import CropService
from superdesk.metadata.item import ITEM_STATE, EMBARGO, SCHEDULE_SETTINGS
from superdesk.utc import utcnow
from superdesk.text_utils import update_word_count
from apps.archive.common import set_sign_off, ITEM_OPERATION, get_user
from apps.archive.archive import flush_renditions
from apps.tasks import send_to
from apps.auth import get_user_id
from .common import BasePublishService, BasePublishResource, ITEM_CORRECT
from superdesk.emails import send_translation_changed
from superdesk.activity import add_activity
from flask import g, current_app as app

def send_translation_notifications(original):
    translated_items = get_resource_service("published").find(
        {"translation_id": original.get("guid"), "last_published_version": True}
    )

    user_ids = set()
    for translated_item in translated_items:
        if translated_item.get("item_id") == original.get("_id"):
            changed_article = translated_item
            continue
        user_ids.add(translated_item.get("original_creator"))
        user_ids.add(translated_item.get("version_creator"))

    if len(user_ids) == 0 or changed_article is None:
        return

    add_activity("translated:changed", "", resource=None, item=changed_article, notify=user_ids)

    recipients = []
    for user_id in user_ids:
        user = get_resource_service("users").find_one(req=None, _id=user_id)
        if user is not None and user.get("email", None) is not None:
            recipients.append(user.get("email"))

    if len(recipients) == 0:
        return

    username = g.user.get("display_name") or g.user.get("username")
    send_translation_changed(username, changed_article, recipients)

class CorrectPublishResource(BasePublishResource):
    def __init__(self, endpoint_name, app, service):
        super().__init__(endpoint_name, app, service, "correct")

class CorrectPublishService(BasePublishService):
    publish_type = "correct"
    published_state = "corrected"
    item_operation = ITEM_CORRECT

    def set_state(self, original, updates):
        updates[ITEM_STATE] = self.published_state

        if original.get(EMBARGO) or updates.get(EMBARGO):
            # embargo time elapsed
            utc_embargo = updates.get(SCHEDULE_SETTINGS, {}).get("utc_{}".format(EMBARGO)) or original.get(
                SCHEDULE_SETTINGS, {}
            ).get("utc_{}".format(EMBARGO))
            if utc_embargo and utc_embargo < utcnow():
                # remove embargo information. so the next correction is without embargo.
                updates[EMBARGO] = None
                super().set_state(original, updates)
        else:
            super().set_state(original, updates)

    def change_being_corrected_to_published(self, updates, original):
        if app.config.get("CORRECTIONS_WORKFLOW") and original.get("state") == "correction":
            publish_service = get_resource_service("published")
            being_corrected_article = publish_service.find_one(
                req=None, guid=original.get("guid"), state="being_corrected"
            )

            if being_corrected_article.get("correction_sequence", 0) > 0:
                publish_service.patch(being_corrected_article["_id"], updates={"state": "corrected"})
            else:
                publish_service.patch(being_corrected_article["_id"], updates={"state": "published"})

    def on_update(self, updates, original):
        CropService().validate_multiple_crops(updates, original)
        super().on_update(updates, original)
        updates[ITEM_OPERATION] = self.item_operation
        updates["versioncreated"] = utcnow()
        updates["correction_sequence"] = original.get("correction_sequence", 1) + 1
        set_sign_off(updates, original)
        update_word_count(updates, original)
        flush_renditions(updates, original)
        self.change_being_corrected_to_published(updates, original)

    def update(self, id, updates, original):
        editor_utils.generate_fields(updates)
        get_resource_service("archive")._handle_media_updates(updates, original, get_user())
        super().update(id, updates, original)

    def on_updated(self, updates, original):
        super().on_updated(updates, original)
        send_translation_notifications(original)
<|endoftext|>