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

import nose
import logging
l = logging.getLogger("angr_tests")

try:
    # pylint: disable=W0611,F0401
    import standard_logging
    import angr_debug
except ImportError:
    pass

import angr

# load the tests
import os
test_location = str(os.path.dirname(os.path.realpath(__file__)))
arger_amd64 = None
arger_i386 = None
arger_ppc32 = None
arger_arm = None
arger_mipsel = None
arger_mips = None

def setup_ppc32():
    global arger_ppc32
    arger_ppc32 = angr.Project(test_location + "/blob/ppc/argc_symbol")

def setup_mips():
    global arger_mips
    arger_mips = angr.Project(test_location + "/blob/mips/argc_symbol")

def setup_mipsel():
    global arger_mipsel
    arger_mipsel = angr.Project(test_location + "/blob/mipsel/argc_symbol")

def setup_amd64():
    global arger_amd64
    arger_amd64 = angr.Project(test_location + "/blob/x86_64/argc_symbol")

def setup_i386():
    global arger_i386
    arger_i386 = angr.Project(test_location + "/blob/i386/argc_symbol")

def setup_arm():
    global arger_arm
    arger_arm = angr.Project(test_location + "/blob/armel/argc_symbol")

def setup_module():
    setup_i386()
    setup_amd64()
    setup_mipsel()
    setup_mips()
    setup_arm()
    setup_ppc32()

def test_mips():
    r_addr = [0x400720, 0x40076c, 0x4007bc]

    s = arger_mips.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_mips, find=r_addr, num_find=100, start=arger_mips.exit_to(arger_mips.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

def test_mipsel():
    r_addr = [0x400720, 0x40076c, 0x4007bc]

    s = arger_mipsel.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_mipsel, find=r_addr, num_find=100, start=arger_mipsel.exit_to(arger_mipsel.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

def test_i386():
    r_addr = [0x08048411, 0x08048437, 0x08048460]

    s = arger_i386.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_i386, find=r_addr, num_find=100, start=arger_i386.exit_to(arger_i386.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

def test_amd64():
    r_addr = [0x40051B, 0x400540, 0x400569]

    s = arger_amd64.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_amd64, find=r_addr, num_find=100, start=arger_amd64.exit_to(arger_amd64.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 800))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 800))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

def test_arm():
    r_addr = [0x00010444, 0x00010478, 0x000104B0]

    s = arger_arm.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_arm, find=r_addr, num_find=100, start=arger_arm.exit_to(arger_arm.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

def test_ppc32():
    r_addr = [0x1000043C, 0x10000474, 0x100004B0]

    s = arger_ppc32.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_ppc32, find=r_addr, num_find=100, start=arger_ppc32.exit_to(arger_ppc32.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

if __name__ == "__main__":
    setup_module()
    #test_mips()
    #test_mipsel()
    #test_arm()
    #test_i386()
    test_amd64()
    #test_ppc32()
<commit_msg>symbolic argc<commit_after>#!/usr/bin/env python

import nose
import logging
l = logging.getLogger("angr_tests")

try:
    # pylint: disable=W0611,F0401
    import standard_logging
    import angr_debug
except ImportError:
    pass

import angr

# load the tests
import os
test_location = str(os.path.dirname(os.path.realpath(__file__)))
arger_amd64 = None
arger_i386 = None
arger_ppc32 = None
arger_arm = None
arger_mipsel = None
arger_mips = None

def setup_ppc32():
    global arger_ppc32
    arger_ppc32 = angr.Project(test_location + "/blob/ppc/argc_symbol")

def setup_mips():
    global arger_mips
    arger_mips = angr.Project(test_location + "/blob/mips/argc_symbol")

def setup_mipsel():
    global arger_mipsel
    arger_mipsel = angr.Project(test_location + "/blob/mipsel/argc_symbol")

def setup_amd64():
    global arger_amd64
    arger_amd64 = angr.Project(test_location + "/blob/x86_64/argc_symbol")

def setup_i386():
    global arger_i386
    arger_i386 = angr.Project(test_location + "/blob/i386/argc_symbol")

def setup_arm():
    global arger_arm
    arger_arm = angr.Project(test_location + "/blob/armel/argc_symbol")

def setup_module():
    setup_i386()
    setup_amd64()
    setup_mipsel()
    setup_mips()
    setup_arm()
    setup_ppc32()

def test_mips():
    r_addr = [0x400720, 0x40076c, 0x4007bc]

    s = arger_mips.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_mips, find=r_addr, num_find=100, start=arger_mips.exit_to(arger_mips.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

def test_mipsel():
    r_addr = [0x400720, 0x40076c, 0x4007bc]

    s = arger_mipsel.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_mipsel, find=r_addr, num_find=100, start=arger_mipsel.exit_to(arger_mipsel.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

def test_i386():
    r_addr = [0x08048411, 0x08048437, 0x08048460]

    s = arger_i386.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_i386, find=r_addr, num_find=100, start=arger_i386.exit_to(arger_i386.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

def test_amd64():
    r_addr = [0x40051B, 0x400540, 0x400569]

    s = arger_amd64.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_amd64, find=r_addr, num_find=100, start=arger_amd64.exit_to(arger_amd64.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 800))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 800))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

def test_arm():
    r_addr = [0x00010444, 0x00010478, 0x000104B0]

    s = arger_arm.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_arm, find=r_addr, num_find=100, start=arger_arm.exit_to(arger_arm.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

def test_ppc32():
    r_addr = [0x1000043C, 0x10000474, 0x100004B0]

    s = arger_ppc32.initial_state(argv = [40, 40, 40], envp ={"HOME": "/home/angr"}, sargc=True)
    xpl = angr.surveyors.Explorer(arger_ppc32, find=r_addr, num_find=100, start=arger_ppc32.exit_to(arger_ppc32.entry, state=s))
    xpl.run()

    nose.tools.assert_equals(len(xpl.found), 3)

    found = xpl.found[0]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals(argc, 0)

    found = xpl.found[1]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Good man" in conc, True)
    nose.tools.assert_equals(argc, 1)

    found = xpl.found[2]
    argc = found._s.se.any_int(found._s['posix'].argc)
    conc = found._s.se.any_str(found._s.mem_expr(found._s.reg_expr('sp'), 400))
    nose.tools.assert_equals("Very Good man" in conc, True)
    nose.tools.assert_equals(argc, 2)

if __name__ == "__main__":
    setup_module()
    test_mips()
    test_mipsel()
    test_arm()
    test_i386()
    test_amd64()
    test_ppc32()
<|endoftext|>


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

import copy
import unittest
import numpy
from pyscf import lib
from pyscf import gto
from pyscf import scf
from pyscf import dft
from pyscf import fci
from pyscf import mcscf

b = 1.4
mol = gto.M(
verbose = 7,
output = '/dev/null',
atom = [
    ['N',(  0.000000,  0.000000, -b/2)],
    ['N',(  0.000000,  0.000000,  b/2)], ],
basis = {'N': '631g', },
)
m = scf.RHF(mol)
m.conv_tol = 1e-10
m.scf()

molsym = gto.M(
verbose = 7,
output = '/dev/null',
atom = [
    ['N',(  0.000000,  0.000000, -b/2)],
    ['N',(  0.000000,  0.000000,  b/2)], ],
basis = {'N': '631g', },
symmetry = True
)
msym = scf.RHF(molsym)
msym.conv_tol = 1e-10
msym.scf()

def tearDownModule():
    global mol, molsym, m, msym
    mol.stdout.close()
    molsym.stdout.close()
    del mol, molsym, m, msym

class KnownValues(unittest.TestCase):
    def test_with_x2c_scanner(self):
        mc1 = mcscf.CASCI(m, 4, 4).x2c().run()
        self.assertAlmostEqual(mc1.e_tot, -108.89264146901512, 7)

        mc1 = mcscf.CASCI(m, 4, 4).x2c().as_scanner().as_scanner()
        mc1(mol)
        self.assertAlmostEqual(mc1.e_tot, -108.89264146901512, 7)

    def test_fix_spin_(self):
        mc1 = mcscf.CASCI(m, 4, 4)
        mc1.fix_spin_(ss=2).run()
        self.assertAlmostEqual(mc1.e_tot, -108.03741684418183, 7)

        mc1.fix_spin_(ss=0).run()
        self.assertAlmostEqual(mc1.e_tot, -108.83741684447352, 7)

    def test_cas_natorb(self):
        mc1 = mcscf.CASCI(msym, 4, 4, ncore=5)
        mo = mc1.sort_mo([4,5,10,13])
        mc1.sorting_mo_energy = True
        mc1.kernel(mo)
        mo0 = mc1.mo_coeff
        ci0 = mc1.ci
        self.assertAlmostEqual(mc1.e_tot, -108.71841138528966, 9)

        mo2, ci2, mo_e = mc1.canonicalize(sort=False, cas_natorb=True,
                                          with_meta_lowdin=False, verbose=7)

        casdm1 = mc1.fcisolver.make_rdm1(mc1.ci, 4, 4)
        mc1.ci = None  # Force cas_natorb_ to recompute CI coefficients
        mc1.cas_natorb_(casdm1=casdm1, sort=True, with_meta_lowdin=True)
        mo1 = mc1.mo_coeff
        ci1 = mc1.ci
        s = numpy.einsum('pi,pq,qj->ij', mo0[:,5:9], msym.get_ovlp(), mo1[:,5:9])
        self.assertAlmostEqual(fci.addons.overlap(ci0, ci1, 4, 4, s), 1, 9)

        self.assertAlmostEqual(float(abs(mo1-mo2).max()), 0, 9)
        self.assertAlmostEqual(ci1.ravel().dot(ci2.ravel()), 1, 9)

    def test_multi_roots(self):
        mc1 = mcscf.CASCI(m, 4, 4)
        mc1.fcisolver.nroots = 2
        mc1.natorb = True
        mc1.kernel()
        self.assertAlmostEqual(mc1.e_tot[0], -108.83741684447352, 9)
        self.assertAlmostEqual(mc1.e_tot[1], -108.72522194135604, 9)
        dm1 = mc1.analyze()
        self.assertAlmostEqual(lib.finger(dm1[0]), 2.6252082970845532, 7)
        self.assertAlmostEqual(lib.finger(dm1[1]), 2.6252082970845532, 7)

    def test_external_fcisolver(self):
        class FCI_as_DMRG(fci.direct_spin1.FCISolver):
            def __getattribute__(self, attr):
                """Prevent 'private' attribute access"""
                if attr in ('make_rdm1s', 'spin_square'):
                    raise AttributeError
                else:
                    return object.__getattribute__(self, attr)
        mc1 = mcscf.CASCI(m, 4, 4)
        mc1.fcisolver = FCI_as_DMRG(mol)
        mc1.fcisolver.nroots = 2
        mc1.natorb = True
        mc1.kernel()
        self.assertAlmostEqual(mc1.e_tot[0], -108.83741684447352, 9)
        self.assertAlmostEqual(mc1.e_tot[1], -108.72522194135604, 9)
        dm1 = mc1.analyze(with_meta_lowdin=False)
        self.assertAlmostEqual(lib.finger(dm1[0]), 2.6252082970845532*2, 7)

    def test_get_h2eff(self):
        mc1 = mcscf.approx_hessian(mcscf.CASCI(m, 4, 4))
        eri1 = mc1.get_h2eff(m.mo_coeff[:,5:9])
        eri2 = mc1.get_h2cas(m.mo_coeff[:,5:9])
        self.assertAlmostEqual(abs(eri1-eri2).max(), 0, 12)

        mc1 = mcscf.density_fit(mcscf.CASCI(m, 4, 4))
        eri1 = mc1.get_h2eff(m.mo_coeff[:,5:9])
        eri2 = mc1.get_h2cas(m.mo_coeff[:,5:9])
        self.assertTrue(abs(eri1-eri2).max() > 1e-5)

    def test_get_veff(self):
        mf = m.view(dft.rks.RKS)
        mc1 = mcscf.CASCI(mf, 4, 4)
        nao = mol.nao_nr()
        dm = numpy.random.random((nao,nao))
        dm = dm + dm.T
        veff1 = mc1.get_veff(mol, dm)
        veff2 = m.get_veff(mol, dm)
        self.assertAlmostEqual(abs(veff1-veff2).max(), 0, 12)

    def test_with_ci_init_guess(self):
        mc1 = mcscf.CASCI(msym, 4, 4)
        ci0 = numpy.zeros((6,6))
        ci0[0,1] = 1
        mc1.kernel(ci0=ci0)

        mc2 = mcscf.CASCI(msym, 4, 4)
        mc2.wfnsym = 'A1u'
        mc2.kernel()
        self.assertAlmostEqual(mc1.e_tot, mc2.e_tot, 9)

    def test_slight_symmetry_broken(self):
        mf = copy.copy(msym)
        mf.mo_coeff = numpy.array(msym.mo_coeff)
        u = numpy.linalg.svd(numpy.random.random((4,4)))[0]
        mf.mo_coeff[:,5:9] = mf.mo_coeff[:,5:9].dot(u)
        mc1 = mcscf.CASCI(mf, 4, 4)
        mc1.kernel()
        self.assertAlmostEqual(mc1.e_tot, -108.83741684445798, 9)

    def test_sort_mo(self):
        mc1 = mcscf.CASCI(msym, 4, 4)
        mo = mc1.sort_mo_by_irrep({'A1u':3, 'A1g':1})
        mc1.kernel(mo)
        self.assertAlmostEqual(mc1.e_tot, -105.82542805259033, 9)

    def test_state_average(self):
        mc = mcscf.CASCI(m, 4, 4)
        mc.state_average_([0.5, 0.25, 0.25])
        mc.fcisolver.spin = 2
        mc.run()
        self.assertAlmostEqual(mc.e_states[0], -108.72522194135607, 9)
        self.assertAlmostEqual(mc.e_states[1], -108.67148843338228, 9)
        self.assertAlmostEqual(mc.e_states[2], -108.67148843338228, 9)

        mc.analyze()
        mo_coeff, civec, mo_occ = mc.cas_natorb(sort=True)

        mc.kernel(mo_coeff=mo_coeff)
        self.assertAlmostEqual(mc.e_states[0], -108.72522194135607, 9)
        self.assertAlmostEqual(mc.e_states[1], -108.67148843338228, 9)
        #FIXME: with the initial guess from mc, FCI solver may converge to
        # another state
        #self.assertAlmostEqual(mc.e_tot[2], -108.67148843338228, 9)
        self.assertAlmostEqual(abs((civec[0]*mc.ci[0]).sum()), 1, 9)
        # Second and third root are degenerated
        #self.assertAlmostEqual(abs((civec[1]*mc.ci[1]).sum()), 1, 9)

    def test_state_average_mix(self):
        mc = mcscf.CASCI(m, 4, 4)
        cis1 = copy.copy(mc.fcisolver)
        cis1.spin = 2
        cis1.nroots = 3
        mc = mcscf.addons.state_average_mix(mc, [cis1, mc.fcisolver], [.25, .25, .25, .25])
        mc.run()
        self.assertAlmostEqual(mc.e_states[0], -108.72522194135607, 9)
        self.assertAlmostEqual(mc.e_states[1], -108.67148843338228, 9)
        self.assertAlmostEqual(mc.e_states[2], -108.67148843338228, 9)
        self.assertAlmostEqual(mc.e_states[3], -108.83741684447352, 9)

        mc.analyze()
        mo_coeff, civec, mo_occ = mc.cas_natorb(sort=True)

        mc.kernel(mo_coeff=mo_coeff)
        self.assertAlmostEqual(mc.e_states[0], -108.72522194135607, 9)
        self.assertAlmostEqual(mc.e_states[1], -108.67148843338228, 9)
        #FIXME: with the initial guess from mc, FCI solver may converge to
        # another state
        #self.assertAlmostEqual(mc.e_states[2], -108.67148843338228, 9)
        self.assertAlmostEqual(mc.e_states[3], -108.83741684447352, 9)
        self.assertAlmostEqual(abs((civec[0]*mc.ci[0]).sum()), 1, 9)
        self.assertAlmostEqual(abs((civec[3]*mc.ci[3]).sum()), 1, 9)

    def test_reset(self):
        myci = mcscf.CASCI(scf.RHF(molsym), 4, 4).density_fit()
        myci.reset(mol)
        self.assertTrue(myci.mol is mol)
        self.assertTrue(myci._scf.mol is mol)
        self.assertTrue(myci.fcisolver.mol is mol)
        self.assertTrue(myci.with_df.mol is mol)

if __name__ == "__main__":
    print("Full Tests for CASCI")
    unittest.main()

<commit_msg>Adding test for mc.mo_occ to test_casci.py<commit_after>#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import unittest
import numpy
from pyscf import lib
from pyscf import gto
from pyscf import scf
from pyscf import dft
from pyscf import fci
from pyscf import mcscf

b = 1.4
mol = gto.M(
verbose = 7,
output = '/dev/null',
atom = [
    ['N',(  0.000000,  0.000000, -b/2)],
    ['N',(  0.000000,  0.000000,  b/2)], ],
basis = {'N': '631g', },
)
m = scf.RHF(mol)
m.conv_tol = 1e-10
m.scf()

molsym = gto.M(
verbose = 7,
output = '/dev/null',
atom = [
    ['N',(  0.000000,  0.000000, -b/2)],
    ['N',(  0.000000,  0.000000,  b/2)], ],
basis = {'N': '631g', },
symmetry = True
)
msym = scf.RHF(molsym)
msym.conv_tol = 1e-10
msym.scf()

def tearDownModule():
    global mol, molsym, m, msym
    mol.stdout.close()
    molsym.stdout.close()
    del mol, molsym, m, msym

class KnownValues(unittest.TestCase):
    def test_with_x2c_scanner(self):
        mc1 = mcscf.CASCI(m, 4, 4).x2c().run()
        self.assertAlmostEqual(mc1.e_tot, -108.89264146901512, 7)

        mc1 = mcscf.CASCI(m, 4, 4).x2c().as_scanner().as_scanner()
        mc1(mol)
        self.assertAlmostEqual(mc1.e_tot, -108.89264146901512, 7)

    def test_fix_spin_(self):
        mc1 = mcscf.CASCI(m, 4, 4)
        mc1.fix_spin_(ss=2).run()
        self.assertAlmostEqual(mc1.e_tot, -108.03741684418183, 7)

        mc1.fix_spin_(ss=0).run()
        self.assertAlmostEqual(mc1.e_tot, -108.83741684447352, 7)

    def test_cas_natorb(self):
        mc1 = mcscf.CASCI(msym, 4, 4, ncore=5)
        mo = mc1.sort_mo([4,5,10,13])
        mc1.sorting_mo_energy = True
        mc1.kernel(mo)
        mo0 = mc1.mo_coeff
        ci0 = mc1.ci
        self.assertAlmostEqual(mc1.e_tot, -108.71841138528966, 9)

        mo2, ci2, mo_e = mc1.canonicalize(sort=False, cas_natorb=True,
                                          with_meta_lowdin=False, verbose=7)

        casdm1 = mc1.fcisolver.make_rdm1(mc1.ci, 4, 4)
        mc1.ci = None  # Force cas_natorb_ to recompute CI coefficients
        mc1.cas_natorb_(casdm1=casdm1, sort=True, with_meta_lowdin=True)
        mo1 = mc1.mo_coeff
        ci1 = mc1.ci
        s = numpy.einsum('pi,pq,qj->ij', mo0[:,5:9], msym.get_ovlp(), mo1[:,5:9])
        self.assertAlmostEqual(fci.addons.overlap(ci0, ci1, 4, 4, s), 1, 9)

        self.assertAlmostEqual(float(abs(mo1-mo2).max()), 0, 9)
        self.assertAlmostEqual(ci1.ravel().dot(ci2.ravel()), 1, 9)

        # Make sure that mc.mo_occ has been set and that the NOONs add to nelectron
        mo_occ = getattr(mc1, "mo_occ", numpy.array([]))
        self.assertNotEqual(mo_occ, numpy.array([]))
        self.assertAlmostEqual(numpy.sum(mo_occ), mc1.mol.nelectron, 9)

    def test_multi_roots(self):
        mc1 = mcscf.CASCI(m, 4, 4)
        mc1.fcisolver.nroots = 2
        mc1.natorb = True
        mc1.kernel()
        self.assertAlmostEqual(mc1.e_tot[0], -108.83741684447352, 9)
        self.assertAlmostEqual(mc1.e_tot[1], -108.72522194135604, 9)
        dm1 = mc1.analyze()
        self.assertAlmostEqual(lib.finger(dm1[0]), 2.6252082970845532, 7)
        self.assertAlmostEqual(lib.finger(dm1[1]), 2.6252082970845532, 7)

    def test_external_fcisolver(self):
        class FCI_as_DMRG(fci.direct_spin1.FCISolver):
            def __getattribute__(self, attr):
                """Prevent 'private' attribute access"""
                if attr in ('make_rdm1s', 'spin_square'):
                    raise AttributeError
                else:
                    return object.__getattribute__(self, attr)
        mc1 = mcscf.CASCI(m, 4, 4)
        mc1.fcisolver = FCI_as_DMRG(mol)
        mc1.fcisolver.nroots = 2
        mc1.natorb = True
        mc1.kernel()
        self.assertAlmostEqual(mc1.e_tot[0], -108.83741684447352, 9)
        self.assertAlmostEqual(mc1.e_tot[1], -108.72522194135604, 9)
        dm1 = mc1.analyze(with_meta_lowdin=False)
        self.assertAlmostEqual(lib.finger(dm1[0]), 2.6252082970845532*2, 7)

    def test_get_h2eff(self):
        mc1 = mcscf.approx_hessian(mcscf.CASCI(m, 4, 4))
        eri1 = mc1.get_h2eff(m.mo_coeff[:,5:9])
        eri2 = mc1.get_h2cas(m.mo_coeff[:,5:9])
        self.assertAlmostEqual(abs(eri1-eri2).max(), 0, 12)

        mc1 = mcscf.density_fit(mcscf.CASCI(m, 4, 4))
        eri1 = mc1.get_h2eff(m.mo_coeff[:,5:9])
        eri2 = mc1.get_h2cas(m.mo_coeff[:,5:9])
        self.assertTrue(abs(eri1-eri2).max() > 1e-5)

    def test_get_veff(self):
        mf = m.view(dft.rks.RKS)
        mc1 = mcscf.CASCI(mf, 4, 4)
        nao = mol.nao_nr()
        dm = numpy.random.random((nao,nao))
        dm = dm + dm.T
        veff1 = mc1.get_veff(mol, dm)
        veff2 = m.get_veff(mol, dm)
        self.assertAlmostEqual(abs(veff1-veff2).max(), 0, 12)

    def test_with_ci_init_guess(self):
        mc1 = mcscf.CASCI(msym, 4, 4)
        ci0 = numpy.zeros((6,6))
        ci0[0,1] = 1
        mc1.kernel(ci0=ci0)

        mc2 = mcscf.CASCI(msym, 4, 4)
        mc2.wfnsym = 'A1u'
        mc2.kernel()
        self.assertAlmostEqual(mc1.e_tot, mc2.e_tot, 9)

    def test_slight_symmetry_broken(self):
        mf = copy.copy(msym)
        mf.mo_coeff = numpy.array(msym.mo_coeff)
        u = numpy.linalg.svd(numpy.random.random((4,4)))[0]
        mf.mo_coeff[:,5:9] = mf.mo_coeff[:,5:9].dot(u)
        mc1 = mcscf.CASCI(mf, 4, 4)
        mc1.kernel()
        self.assertAlmostEqual(mc1.e_tot, -108.83741684445798, 9)

    def test_sort_mo(self):
        mc1 = mcscf.CASCI(msym, 4, 4)
        mo = mc1.sort_mo_by_irrep({'A1u':3, 'A1g':1})
        mc1.kernel(mo)
        self.assertAlmostEqual(mc1.e_tot, -105.82542805259033, 9)

    def test_state_average(self):
        mc = mcscf.CASCI(m, 4, 4)
        mc.state_average_([0.5, 0.25, 0.25])
        mc.fcisolver.spin = 2
        mc.run()
        self.assertAlmostEqual(mc.e_states[0], -108.72522194135607, 9)
        self.assertAlmostEqual(mc.e_states[1], -108.67148843338228, 9)
        self.assertAlmostEqual(mc.e_states[2], -108.67148843338228, 9)

        mc.analyze()
        mo_coeff, civec, mo_occ = mc.cas_natorb(sort=True)

        mc.kernel(mo_coeff=mo_coeff)
        self.assertAlmostEqual(mc.e_states[0], -108.72522194135607, 9)
        self.assertAlmostEqual(mc.e_states[1], -108.67148843338228, 9)
        #FIXME: with the initial guess from mc, FCI solver may converge to
        # another state
        #self.assertAlmostEqual(mc.e_tot[2], -108.67148843338228, 9)
        self.assertAlmostEqual(abs((civec[0]*mc.ci[0]).sum()), 1, 9)
        # Second and third root are degenerated
        #self.assertAlmostEqual(abs((civec[1]*mc.ci[1]).sum()), 1, 9)

    def test_state_average_mix(self):
        mc = mcscf.CASCI(m, 4, 4)
        cis1 = copy.copy(mc.fcisolver)
        cis1.spin = 2
        cis1.nroots = 3
        mc = mcscf.addons.state_average_mix(mc, [cis1, mc.fcisolver], [.25, .25, .25, .25])
        mc.run()
        self.assertAlmostEqual(mc.e_states[0], -108.72522194135607, 9)
        self.assertAlmostEqual(mc.e_states[1], -108.67148843338228, 9)
        self.assertAlmostEqual(mc.e_states[2], -108.67148843338228, 9)
        self.assertAlmostEqual(mc.e_states[3], -108.83741684447352, 9)

        mc.analyze()
        mo_coeff, civec, mo_occ = mc.cas_natorb(sort=True)

        mc.kernel(mo_coeff=mo_coeff)
        self.assertAlmostEqual(mc.e_states[0], -108.72522194135607, 9)
        self.assertAlmostEqual(mc.e_states[1], -108.67148843338228, 9)
        #FIXME: with the initial guess from mc, FCI solver may converge to
        # another state
        #self.assertAlmostEqual(mc.e_states[2], -108.67148843338228, 9)
        self.assertAlmostEqual(mc.e_states[3], -108.83741684447352, 9)
        self.assertAlmostEqual(abs((civec[0]*mc.ci[0]).sum()), 1, 9)
        self.assertAlmostEqual(abs((civec[3]*mc.ci[3]).sum()), 1, 9)

    def test_reset(self):
        myci = mcscf.CASCI(scf.RHF(molsym), 4, 4).density_fit()
        myci.reset(mol)
        self.assertTrue(myci.mol is mol)
        self.assertTrue(myci._scf.mol is mol)
        self.assertTrue(myci.fcisolver.mol is mol)
        self.assertTrue(myci.with_df.mol is mol)

if __name__ == "__main__":
    print("Full Tests for CASCI")
    unittest.main()

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Test bulk copy.
"""

import unittest
import datetime

from pymssql import _mssql
from tests.helpers import drop_table, pymssqlconn, mssql_server_required

tablename = "pymssql"
simple_table = "CREATE TABLE %s (a1 INT, a2 INT, a3 INT)" % tablename
complex_table = """
    CREATE TABLE %s (
        pk_id int IDENTITY (1, 1) NOT NULL,
        uuid uniqueidentifier DEFAULT newsequentialid(),
        col_real real UNIQUE,
        col_float float,
        col_datetime datetime,
        col_bit bit,
        col_varchar varchar(50)
    )
""" % tablename

@mssql_server_required
class TestTypes(unittest.TestCase):
    def setUp(self):
        self.conn = pymssqlconn()
        drop_table(self.conn._conn, tablename)

    def tearDown(self):
        self.conn.close()

    def expect_simple_table_content(self, query, content):
        self.conn._conn.execute_query(query)
        assert [(row[0], row[1], row[2]) for row in self.conn._conn] == content

    def expect_row_count(self, expected_row_count):
        self.conn._conn.execute_query('select count(*) from pymssql')
        assert tuple(self.conn._conn)[0][0] == expected_row_count

    def simple_table_test(self, content, **kwargs):
        self.conn._conn.execute_non_query(simple_table)
        self.conn.bulk_copy(tablename, content, **kwargs)
        self.expect_simple_table_content('select * from pymssql', content)

    def test_simple_table_bulk_copy(self):
        self.simple_table_test([(1, 2, 3), (4, 5, 6)])

    def test_lots_of_rows_single_batch(self):
        self.conn._conn.execute_non_query(simple_table)
        self.conn.bulk_copy(tablename, [(1, 2, 3), (4, 5, 6)] * 100000, batch_size=1000000)
        self.expect_simple_table_content('select top 2 * from pymssql', [(1, 2, 3), (4, 5, 6)])
        self.expect_row_count(200000)

    def test_batches(self):
        self.conn._conn.execute_non_query(simple_table)

        self.conn.bulk_copy(tablename, [(1, 2, 3), (4, 5, 6)] * 100000, batch_size=1000)

        self.expect_simple_table_content('select top 2 * from pymssql', [(1, 2, 3), (4, 5, 6)])
        self.expect_row_count(200000)

    def test_exact_batch_size(self):
        self.conn._conn.execute_non_query(simple_table)

        self.conn.bulk_copy(tablename, [(1, 2, 3), (4, 5, 6)] * 500, batch_size=1000)

        self.expect_simple_table_content('select top 2 * from pymssql', [(1, 2, 3), (4, 5, 6)])
        self.expect_row_count(1000)

    def test_tablock_hint(self):
        self.simple_table_test([(1, 2, 3), (4, 5, 6)], tablock=True)

    def test_check_constraints_hint(self):
        self.simple_table_test([(1, 2, 3), (4, 5, 6)], check_constraints=True)

    def test_fire_triggers_hint(self):
        self.simple_table_test([(1, 2, 3), (4, 5, 6)], fire_triggers=True)

    def test_null_values(self):
        self.simple_table_test([(1, None, 3), (None, None, None), (1, 2, 3)])

    def test_column_ids(self):
        self.conn._conn.execute_non_query(simple_table)
        self.conn.bulk_copy(tablename, [(1, 2, 3), (4, 5, 6)], column_ids=[1, 3, 2])
        self.expect_simple_table_content('select * from pymssql', [(1, 3, 2), (4, 6, 5)])

    def test_too_many_columns(self):
        self.conn._conn.execute_non_query(simple_table)
        caught_exception = False

        try:
            self.conn.bulk_copy(tablename, [(7, 7, 7, 7)])
        except _mssql.MSSQLDatabaseException:
            caught_exception = True

        assert caught_exception

    def test_bad_value(self):
        self.conn._conn.execute_non_query(simple_table)
        caught_exception = False

        try:
            self.conn.bulk_copy(tablename, [("Hello", 7, 7)])
        except _mssql.MSSQLDatabaseException:
            caught_exception = True

        assert caught_exception

    def test_too_few_column_ids(self):
        self.conn._conn.execute_non_query(simple_table)
        caught_exception = False

        try:
            self.conn.bulk_copy(tablename, [(1, 2, 3)], column_ids=[1])
        except ValueError:
            caught_exception = True

        assert caught_exception

    def test_invalid_column_ids(self):
        self.conn._conn.execute_non_query(simple_table)
        caught_exception = False

        try:
            self.conn.bulk_copy(tablename, [(1, 2, 3)], column_ids=[1, 2, 9])
        except _mssql.MSSQLDatabaseException:
            caught_exception = True

        assert caught_exception

    def test_complex_table(self):
        self.conn._conn.execute_non_query(complex_table)
        rows = [
            (1.2000000476837158, 3.4, datetime.datetime(year=2020, month=1, day=2, hour=3, minute=4, second=5), True, "Hello World"),
            (5.599999904632568, 7.8, datetime.datetime(year=2021, month=2, day=3, hour=4, minute=5, second=6), False, "Hello World!"),
        ]
        self.conn.bulk_copy(tablename, rows, [3, 4, 5, 6, 7])
        self.conn._conn.execute_query('select * from pymssql')
        assert [tuple(row[i] for i in range(2, 7)) for row in self.conn._conn] == rows

    def test_uniqueness_failure(self):
        self.conn._conn.execute_non_query(complex_table)
        caught_exception = False

        rows = [
            (1.2000000476837158, 3.4, datetime.datetime(year=2020, month=1, day=2, hour=3, minute=4, second=5), True, "Hello World"),
            (1.2000000476837158, 7.8, datetime.datetime(year=2021, month=2, day=3, hour=4, minute=5, second=6), False, "Hello World!"),
        ]
        try:
            self.conn.bulk_copy(tablename, rows, [3, 4, 5, 6, 7])
        except _mssql.MSSQLDatabaseException:
            caught_exception = True

        assert caught_exception
<commit_msg>Use assertRaises context manager in test_bulk_copy.py.<commit_after># -*- coding: utf-8 -*-
"""
Test bulk copy.
"""

import unittest
import datetime

from pymssql import _mssql
from tests.helpers import drop_table, pymssqlconn, mssql_server_required

tablename = "pymssql"
simple_table = "CREATE TABLE %s (a1 INT, a2 INT, a3 INT)" % tablename
complex_table = """
    CREATE TABLE %s (
        pk_id int IDENTITY (1, 1) NOT NULL,
        uuid uniqueidentifier DEFAULT newsequentialid(),
        col_real real UNIQUE,
        col_float float,
        col_datetime datetime,
        col_bit bit,
        col_varchar varchar(50)
    )
""" % tablename

@mssql_server_required
class TestTypes(unittest.TestCase):
    def setUp(self):
        self.conn = pymssqlconn()
        drop_table(self.conn._conn, tablename)

    def tearDown(self):
        self.conn.close()

    def expect_simple_table_content(self, query, content):
        self.conn._conn.execute_query(query)
        assert [(row[0], row[1], row[2]) for row in self.conn._conn] == content

    def expect_row_count(self, expected_row_count):
        self.conn._conn.execute_query('select count(*) from pymssql')
        assert tuple(self.conn._conn)[0][0] == expected_row_count

    def simple_table_test(self, content, **kwargs):
        self.conn._conn.execute_non_query(simple_table)
        self.conn.bulk_copy(tablename, content, **kwargs)
        self.expect_simple_table_content('select * from pymssql', content)

    def test_simple_table_bulk_copy(self):
        self.simple_table_test([(1, 2, 3), (4, 5, 6)])

    def test_lots_of_rows_single_batch(self):
        self.conn._conn.execute_non_query(simple_table)
        self.conn.bulk_copy(tablename, [(1, 2, 3), (4, 5, 6)] * 100000, batch_size=1000000)
        self.expect_simple_table_content('select top 2 * from pymssql', [(1, 2, 3), (4, 5, 6)])
        self.expect_row_count(200000)

    def test_batches(self):
        self.conn._conn.execute_non_query(simple_table)

        self.conn.bulk_copy(tablename, [(1, 2, 3), (4, 5, 6)] * 100000, batch_size=1000)

        self.expect_simple_table_content('select top 2 * from pymssql', [(1, 2, 3), (4, 5, 6)])
        self.expect_row_count(200000)

    def test_exact_batch_size(self):
        self.conn._conn.execute_non_query(simple_table)

        self.conn.bulk_copy(tablename, [(1, 2, 3), (4, 5, 6)] * 500, batch_size=1000)

        self.expect_simple_table_content('select top 2 * from pymssql', [(1, 2, 3), (4, 5, 6)])
        self.expect_row_count(1000)

    def test_tablock_hint(self):
        self.simple_table_test([(1, 2, 3), (4, 5, 6)], tablock=True)

    def test_check_constraints_hint(self):
        self.simple_table_test([(1, 2, 3), (4, 5, 6)], check_constraints=True)

    def test_fire_triggers_hint(self):
        self.simple_table_test([(1, 2, 3), (4, 5, 6)], fire_triggers=True)

    def test_null_values(self):
        self.simple_table_test([(1, None, 3), (None, None, None), (1, 2, 3)])

    def test_column_ids(self):
        self.conn._conn.execute_non_query(simple_table)
        self.conn.bulk_copy(tablename, [(1, 2, 3), (4, 5, 6)], column_ids=[1, 3, 2])
        self.expect_simple_table_content('select * from pymssql', [(1, 3, 2), (4, 6, 5)])

    def test_too_many_columns(self):
        self.conn._conn.execute_non_query(simple_table)
        with self.assertRaises(_mssql.MSSQLDatabaseException):
            self.conn.bulk_copy(tablename, [(7, 7, 7, 7)])

    def test_bad_value(self):
        self.conn._conn.execute_non_query(simple_table)
        with self.assertRaises(_mssql.MSSQLDatabaseException):
            self.conn.bulk_copy(tablename, [("Hello", 7, 7)])

    def test_too_few_column_ids(self):
        self.conn._conn.execute_non_query(simple_table)
        caught_exception = False

        try:
            self.conn.bulk_copy(tablename, [(1, 2, 3)], column_ids=[1])
        except ValueError:
            caught_exception = True

        assert caught_exception

    def test_invalid_column_ids(self):
        self.conn._conn.execute_non_query(simple_table)
        with self.assertRaises(_mssql.MSSQLDatabaseException):
            self.conn.bulk_copy(tablename, [(1, 2, 3)], column_ids=[1, 2, 9])

    def test_complex_table(self):
        self.conn._conn.execute_non_query(complex_table)
        rows = [
            (1.2000000476837158, 3.4, datetime.datetime(year=2020, month=1, day=2, hour=3, minute=4, second=5), True, "Hello World"),
            (5.599999904632568, 7.8, datetime.datetime(year=2021, month=2, day=3, hour=4, minute=5, second=6), False, "Hello World!"),
        ]
        self.conn.bulk_copy(tablename, rows, [3, 4, 5, 6, 7])
        self.conn._conn.execute_query('select * from pymssql')
        assert [tuple(row[i] for i in range(2, 7)) for row in self.conn._conn] == rows

    def test_uniqueness_failure(self):
        self.conn._conn.execute_non_query(complex_table)

        rows = [
            (1.2000000476837158, 3.4, datetime.datetime(year=2020, month=1, day=2, hour=3, minute=4, second=5), True, "Hello World"),
            (1.2000000476837158, 7.8, datetime.datetime(year=2021, month=2, day=3, hour=4, minute=5, second=6), False, "Hello World!"),
        ]
        with self.assertRaises(_mssql.MSSQLDatabaseException):
            self.conn.bulk_copy(tablename, rows, [3, 4, 5, 6, 7])
<|endoftext|>


<commit_before>#
# Copyright (c) 2005-2006 rPath, Inc.
#
# All Rights Reserved
#
import base64
import os
import sys
import traceback

from mod_python import apache

from conary.server import http
from conary.repository import errors
from conary.repository.shimclient import ShimNetClient
from conary import conaryclient

from webhandler import WebHandler, normPath, HttpForbidden, HttpNotFound
from templates import repos
from mint.session import SqlSession
from mint import database

class ConaryHandler(WebHandler, http.HttpHandler):
    def __init__(self, req, cfg, repServer = None):
        protocol = 'http'
        port = 80

        #If the browser can support it, give it what it wants.
        if 'application/xhtml+xml' in req.headers_in.get('Accept', ''):
            self.content_type = 'application/xhtml+xml'
            self.output = 'xhtml'

        if repServer:
            self.repServer = repServer
            self.troveStore = self.repServer.troveStore
        if 'mint.web.templates.repos' in sys.modules:
            self.reposTemplatePath = os.path.dirname(sys.modules['mint.web.templates.repos'].__file__) + "/repos/"

    def handle(self, context):
        self.__dict__.update(**context)

        # set up the netclient
        self.serverName = self.req.hostname

        path = self.req.path_info.split("/")
        self.cmd = path[3]
        try:
            if path[1] == "repos":
                self.project = self.client.getProjectByHostname(path[2])
                self.serverName = self.project.getLabel().split("@")[0]
            else:
                self.project = self.client.getProjectByFQDN(self.serverName)
        except database.ItemNotFound:
            raise HttpNotFound

        self.basePath += "repos/%s" % self.project.getHostname()
        self.basePath = normPath(self.basePath)

        if self.auth:
            self.userLevel = self.project.getUserLevel(self.auth.userId)

        return self._handle

    def _handle(self, *args, **kwargs):
        """Handle either an HTTP POST or GET command."""

        if self.project.external:
            useSSL = None # use the label as-is for external projects
        else:
            useSSL = self.cfg.SSL
        cfg = self.project.getConaryConfig(overrideSSL = True, overrideAuth = True,
                                           newUser = self.authToken[0],
                                           newPass = self.authToken[1],
                                           useSSL = useSSL)
        self.authToken = (self.authToken[0], self.authToken[1], None, None)

        if self.project.external:
            self.repos = conaryclient.ConaryClient(cfg).getRepos()
        else:
            self.repos = ShimNetClient(self.repServer, 'http', 80,
                                       self.authToken, cfg.repositoryMap,
                                       cfg.user)

        try:
            method = self.__getattribute__(self.cmd)
        except AttributeError:
            return apache.HTTP_NOT_FOUND

        d = self.fields
        d['auth'] = self.authToken
        try:
            output = method(**d)
            return output
        except errors.OpenError:
            return self._requestAuth()

    def _requestAuth(self):
        # try to fall back to anonymous and rerun the handler
        if self.authToken != ('anonymous', 'anonymous'):
            self.authToken = ('anonymous', 'anonymous')
            return self._handle()
        else:
            raise HttpForbidden

    def _write(self, templateName, **values):
        return WebHandler._write(self, templateName, templatePath = self.reposTemplatePath, **values)

    # XXX: I'm wondering if this is the best approach.
    # maybe it would be better to override _getHandler and
    # provide a list of explicitly-allowed methods. opt-in
    # instead of opt-out.
    #
    # remove some disallowed methods for mint:
    del http.HttpHandler.main
    del http.HttpHandler.metadata
    del http.HttpHandler.chooseBranch
    del http.HttpHandler.getMetadata
    del http.HttpHandler.updateMetadata
    del http.HttpHandler.userlist
    del http.HttpHandler.deleteGroup
    del http.HttpHandler.deleteUser
    del http.HttpHandler.addPermForm
    del http.HttpHandler.addPerm
    del http.HttpHandler.addGroupForm
    del http.HttpHandler.manageGroupForm
    del http.HttpHandler.manageGroup
    del http.HttpHandler.addGroup
    del http.HttpHandler.deletePerm
    del http.HttpHandler.addUser
    del http.HttpHandler.chPassForm
    del http.HttpHandler.chPass
<commit_msg>fix a possible bug since we modify self.authToken later<commit_after>#
# Copyright (c) 2005-2006 rPath, Inc.
#
# All Rights Reserved
#
import base64
import os
import sys
import traceback

from mod_python import apache

from conary.server import http
from conary.repository import errors
from conary.repository.shimclient import ShimNetClient
from conary import conaryclient

from webhandler import WebHandler, normPath, HttpForbidden, HttpNotFound
from templates import repos
from mint.session import SqlSession
from mint import database

class ConaryHandler(WebHandler, http.HttpHandler):
    def __init__(self, req, cfg, repServer = None):
        protocol = 'http'
        port = 80

        #If the browser can support it, give it what it wants.
        if 'application/xhtml+xml' in req.headers_in.get('Accept', ''):
            self.content_type = 'application/xhtml+xml'
            self.output = 'xhtml'

        if repServer:
            self.repServer = repServer
            self.troveStore = self.repServer.troveStore
        if 'mint.web.templates.repos' in sys.modules:
            self.reposTemplatePath = os.path.dirname(sys.modules['mint.web.templates.repos'].__file__) + "/repos/"

    def handle(self, context):
        self.__dict__.update(**context)

        # set up the netclient
        self.serverName = self.req.hostname

        path = self.req.path_info.split("/")
        self.cmd = path[3]
        try:
            if path[1] == "repos":
                self.project = self.client.getProjectByHostname(path[2])
                self.serverName = self.project.getLabel().split("@")[0]
            else:
                self.project = self.client.getProjectByFQDN(self.serverName)
        except database.ItemNotFound:
            raise HttpNotFound

        self.basePath += "repos/%s" % self.project.getHostname()
        self.basePath = normPath(self.basePath)

        if self.auth:
            self.userLevel = self.project.getUserLevel(self.auth.userId)

        return self._handle

    def _handle(self, *args, **kwargs):
        """Handle either an HTTP POST or GET command."""

        if self.project.external:
            useSSL = None # use the label as-is for external projects
        else:
            useSSL = self.cfg.SSL
        cfg = self.project.getConaryConfig(overrideSSL = True, overrideAuth = True,
                                           newUser = self.authToken[0],
                                           newPass = self.authToken[1],
                                           useSSL = useSSL)
        self.authToken = (self.authToken[0], self.authToken[1], None, None)

        if self.project.external:
            self.repos = conaryclient.ConaryClient(cfg).getRepos()
        else:
            self.repos = ShimNetClient(self.repServer, 'http', 80,
                                       self.authToken, cfg.repositoryMap,
                                       cfg.user)

        try:
            method = self.__getattribute__(self.cmd)
        except AttributeError:
            return apache.HTTP_NOT_FOUND

        d = self.fields
        d['auth'] = self.authToken
        try:
            output = method(**d)
            return output
        except errors.OpenError:
            return self._requestAuth()

    def _requestAuth(self):
        # try to fall back to anonymous and rerun the handler
        if self.authToken[0] != 'anonymous':
            self.authToken = ('anonymous', 'anonymous')
            return self._handle()
        else:
            raise HttpForbidden

    def _write(self, templateName, **values):
        return WebHandler._write(self, templateName, templatePath = self.reposTemplatePath, **values)

    # XXX: I'm wondering if this is the best approach.
    # maybe it would be better to override _getHandler and
    # provide a list of explicitly-allowed methods. opt-in
    # instead of opt-out.
    #
    # remove some disallowed methods for mint:
    del http.HttpHandler.main
    del http.HttpHandler.metadata
    del http.HttpHandler.chooseBranch
    del http.HttpHandler.getMetadata
    del http.HttpHandler.updateMetadata
    del http.HttpHandler.userlist
    del http.HttpHandler.deleteGroup
    del http.HttpHandler.deleteUser
    del http.HttpHandler.addPermForm
    del http.HttpHandler.addPerm
    del http.HttpHandler.addGroupForm
    del http.HttpHandler.manageGroupForm
    del http.HttpHandler.manageGroup
    del http.HttpHandler.addGroup
    del http.HttpHandler.deletePerm
    del http.HttpHandler.addUser
    del http.HttpHandler.chPassForm
    del http.HttpHandler.chPass
<|endoftext|>


<commit_before>#-*- coding: utf-8 -*-
from django import forms
from django.contrib.admin import widgets
from django.contrib.auth.decorators import login_required
from django.core.exceptions import PermissionDenied
from django.http import HttpResponseRedirect, HttpResponse
from django.shortcuts import render_to_response
from django.template import RequestContext
from django.utils.translation import ugettext_lazy as _
from models import Folder, Image, Clipboard, tools, FolderRoot
from filer import settings as filer_settings

class NewFolderForm(forms.ModelForm):
    class Meta:
        model = Folder
        fields = ('name',)
        widgets = {
            'name': widgets.AdminTextInputWidget,
        }

def popup_status(request):
    return ('_popup' in request.REQUEST or 'pop' in request.REQUEST)

def selectfolder_status(request):
    return ('select_folder' in request.REQUEST)

def popup_param(request, separator="?"):
    if popup_status(request):
        return "%s_popup=1" % separator
    else:
        return ""

def selectfolder_param(request, separator="&"):
    if selectfolder_status(request):
        return "%sselect_folder=1" % separator
    else:
        return ""

def _userperms(item, request):
    r = []
    ps = ['read', 'edit', 'add_children']
    for p in ps:
        attr = "has_%s_permission" % p
        if hasattr(item, attr):
            x = getattr(item, attr)(request)
            if x:
                r.append(p)
    return r

@login_required
def edit_folder(request, folder_id):
    # TODO: implement edit_folder view
    folder = None
    return render_to_response('admin/filer/folder/folder_edit.html', {
            'folder': folder,
            'is_popup': popup_status(request),
            'select_folder': selectfolder_status(request),
        }, context_instance=RequestContext(request))

@login_required
def edit_image(request, folder_id):
    # TODO: implement edit_image view
    folder = None
    return render_to_response('filer/image_edit.html', {
            'folder': folder,
            'is_popup': popup_status(request),
            'select_folder': selectfolder_status(request),
        }, context_instance=RequestContext(request))

@login_required
def make_folder(request, folder_id=None):
    if not folder_id:
        folder_id = request.REQUEST.get('parent_id', None)
    if folder_id:
        folder = Folder.objects.get(id=folder_id)
    else:
        folder = None

    if request.user.is_superuser:
        pass
    elif folder == None:
        # regular users may not add root folders unless configured otherwise
        if not filer_settings.FILER_ALLOW_REGULAR_USERS_TO_ADD_ROOT_FOLDERS:
            raise PermissionDenied
    elif not folder.has_add_children_permission(request):
        # the user does not have the permission to add subfolders
        raise PermissionDenied

    if request.method == 'POST':
        new_folder_form = NewFolderForm(request.POST)
        if new_folder_form.is_valid():
            new_folder = new_folder_form.save(commit=False)
            if (folder or FolderRoot()).contains_folder(new_folder.name):
                new_folder_form._errors['name'] = new_folder_form.error_class([_('Folder with this name already exists.')])
            else:
                new_folder.parent = folder
                new_folder.owner = request.user
                new_folder.save()
                return HttpResponse('')
    else:
        new_folder_form = NewFolderForm()
    return render_to_response('admin/filer/folder/new_folder_form.html', {
            'new_folder_form': new_folder_form,
            'is_popup': popup_status(request),
            'select_folder': selectfolder_status(request),
    }, context_instance=RequestContext(request))

class UploadFileForm(forms.ModelForm):
    class Meta:
        model = Image

@login_required
def upload(request):
    return render_to_response('filer/upload.html', {
                    'title': u'Upload files',
                    'is_popup': popup_status(request),
                    'select_folder': selectfolder_status(request),
                    }, context_instance=RequestContext(request))

@login_required
def paste_clipboard_to_folder(request):
    if request.method == 'POST':
        folder = Folder.objects.get(id=request.POST.get('folder_id'))
        clipboard = Clipboard.objects.get(id=request.POST.get('clipboard_id'))
        if folder.has_add_children_permission(request):
            tools.move_files_from_clipboard_to_folder(clipboard, folder)
            tools.discard_clipboard(clipboard)
        else:
            raise PermissionDenied
    return HttpResponseRedirect('%s%s%s' % (
                                    request.REQUEST.get('redirect_to', ''),
                                    popup_param(request),
                                    selectfolder_param(request)))

@login_required
def discard_clipboard(request):
    if request.method == 'POST':
        clipboard = Clipboard.objects.get(id=request.POST.get('clipboard_id'))
        tools.discard_clipboard(clipboard)
    return HttpResponseRedirect('%s%s%s' % (
                                    request.POST.get('redirect_to', ''),
                                    popup_param(request),
                                    selectfolder_param(request)))

@login_required
def delete_clipboard(request):
    if request.method == 'POST':
        clipboard = Clipboard.objects.get(id=request.POST.get('clipboard_id'))
        tools.delete_clipboard(clipboard)
    return HttpResponseRedirect('%s%s%s' % (
                                    request.POST.get('redirect_to', ''),
                                    popup_param(request),
                                    selectfolder_param(request)))

@login_required
def clone_files_from_clipboard_to_folder(request):
    if request.method == 'POST':
        clipboard = Clipboard.objects.get(id=request.POST.get('clipboard_id'))
        folder = Folder.objects.get(id=request.POST.get('folder_id'))
        tools.clone_files_from_clipboard_to_folder(clipboard, folder)
    return HttpResponseRedirect('%s%s%s' % (
                                    request.POST.get('redirect_to', ''),
                                    popup_param(request),
                                    selectfolder_param(request)))
<commit_msg>Add some extra markup to the make_folder response to combat issues where the response is being cut off at the end.<commit_after>#-*- coding: utf-8 -*-
from django import forms
from django.contrib.admin import widgets
from django.contrib.auth.decorators import login_required
from django.core.exceptions import PermissionDenied
from django.http import HttpResponseRedirect, HttpResponse
from django.shortcuts import render_to_response
from django.template import RequestContext
from django.utils.translation import ugettext_lazy as _
from models import Folder, Image, Clipboard, tools, FolderRoot
from filer import settings as filer_settings

class NewFolderForm(forms.ModelForm):
    class Meta:
        model = Folder
        fields = ('name',)
        widgets = {
            'name': widgets.AdminTextInputWidget,
        }

def popup_status(request):
    return ('_popup' in request.REQUEST or 'pop' in request.REQUEST)

def selectfolder_status(request):
    return ('select_folder' in request.REQUEST)

def popup_param(request, separator="?"):
    if popup_status(request):
        return "%s_popup=1" % separator
    else:
        return ""

def selectfolder_param(request, separator="&"):
    if selectfolder_status(request):
        return "%sselect_folder=1" % separator
    else:
        return ""

def _userperms(item, request):
    r = []
    ps = ['read', 'edit', 'add_children']
    for p in ps:
        attr = "has_%s_permission" % p
        if hasattr(item, attr):
            x = getattr(item, attr)(request)
            if x:
                r.append(p)
    return r

@login_required
def edit_folder(request, folder_id):
    # TODO: implement edit_folder view
    folder = None
    return render_to_response('admin/filer/folder/folder_edit.html', {
            'folder': folder,
            'is_popup': popup_status(request),
            'select_folder': selectfolder_status(request),
        }, context_instance=RequestContext(request))

@login_required
def edit_image(request, folder_id):
    # TODO: implement edit_image view
    folder = None
    return render_to_response('filer/image_edit.html', {
            'folder': folder,
            'is_popup': popup_status(request),
            'select_folder': selectfolder_status(request),
        }, context_instance=RequestContext(request))

@login_required
def make_folder(request, folder_id=None):
    if not folder_id:
        folder_id = request.REQUEST.get('parent_id', None)
    if folder_id:
        folder = Folder.objects.get(id=folder_id)
    else:
        folder = None

    if request.user.is_superuser:
        pass
    elif folder == None:
        # regular users may not add root folders unless configured otherwise
        if not filer_settings.FILER_ALLOW_REGULAR_USERS_TO_ADD_ROOT_FOLDERS:
            raise PermissionDenied
    elif not folder.has_add_children_permission(request):
        # the user does not have the permission to add subfolders
        raise PermissionDenied

    if request.method == 'POST':
        new_folder_form = NewFolderForm(request.POST)
        if new_folder_form.is_valid():
            new_folder = new_folder_form.save(commit=False)
            if (folder or FolderRoot()).contains_folder(new_folder.name):
                new_folder_form._errors['name'] = new_folder_form.error_class([_('Folder with this name already exists.')])
            else:
                new_folder.parent = folder
                new_folder.owner = request.user
                new_folder.save()
                return HttpResponse('')
    else:
        new_folder_form = NewFolderForm()
    return render_to_response('admin/filer/folder/new_folder_form.html', {
            'new_folder_form': new_folder_form,
            'is_popup': popup_status(request),
            'select_folder': selectfolder_status(request),
    }, context_instance=RequestContext(request))

class UploadFileForm(forms.ModelForm):
    class Meta:
        model = Image

@login_required
def upload(request):
    return render_to_response('filer/upload.html', {
                    'title': u'Upload files',
                    'is_popup': popup_status(request),
                    'select_folder': selectfolder_status(request),
                    }, context_instance=RequestContext(request))

@login_required
def paste_clipboard_to_folder(request):
    if request.method == 'POST':
        folder = Folder.objects.get(id=request.POST.get('folder_id'))
        clipboard = Clipboard.objects.get(id=request.POST.get('clipboard_id'))
        if folder.has_add_children_permission(request):
            tools.move_files_from_clipboard_to_folder(clipboard, folder)
            tools.discard_clipboard(clipboard)
        else:
            raise PermissionDenied
    return HttpResponseRedirect('%s%s%s' % (
                                    request.REQUEST.get('redirect_to', ''),
                                    popup_param(request),
                                    selectfolder_param(request)))

@login_required
def discard_clipboard(request):
    if request.method == 'POST':
        clipboard = Clipboard.objects.get(id=request.POST.get('clipboard_id'))
        tools.discard_clipboard(clipboard)
    return HttpResponseRedirect('%s%s%s' % (
                                    request.POST.get('redirect_to', ''),
                                    popup_param(request),
                                    selectfolder_param(request)))

@login_required
def delete_clipboard(request):
    if request.method == 'POST':
        clipboard = Clipboard.objects.get(id=request.POST.get('clipboard_id'))
        tools.delete_clipboard(clipboard)
    return HttpResponseRedirect('%s%s%s' % (
                                    request.POST.get('redirect_to', ''),
                                    popup_param(request),
                                    selectfolder_param(request)))

@login_required
def clone_files_from_clipboard_to_folder(request):
    if request.method == 'POST':
        clipboard = Clipboard.objects.get(id=request.POST.get('clipboard_id'))
        folder = Folder.objects.get(id=request.POST.get('folder_id'))
        tools.clone_files_from_clipboard_to_folder(clipboard, folder)
    return HttpResponseRedirect('%s%s%s' % (
                                    request.POST.get('redirect_to', ''),
                                    popup_param(request),
                                    selectfolder_param(request)))
<|endoftext|>


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

#   This file is part of periscope.
#   Copyright (c) 2008-2011 Patrick Dessalle <patrick@dessalle.be>
#
#    periscope is free software; you can redistribute it and/or modify
#    it under the terms of the GNU Lesser General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    periscope is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public License
#    along with periscope; if not, write to the Free Software
#    Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

import os, struct, xmlrpclib, commands, gzip, traceback, logging
import xbmc
import socket # For timeout purposes

import SubtitleDatabase

log = logging.getLogger(__name__)

OS_LANGS ={ "en": "eng",
            "fr" : "fre",
            "hu": "hun",
            "cs": "cze",
            "pl" : "pol",
            "sk" : "slo",
            "pt" : "por",
            "pt-br" : "pob",
            "es" : "spa",
            "el" : "ell",
            "ar":"ara",
            'sq':'alb',
            "hy":"arm",
            "ay":"ass",
            "bs":"bos",
            "bg":"bul",
            "ca":"cat",
            "zh":"chi",
            "hr":"hrv",
            "da":"dan",
            "nl":"dut",
            "eo":"epo",
            "et":"est",
            "fi":"fin",
            "gl":"glg",
            "ka":"geo",
            "de":"ger",
            "he":"heb",
            "hi":"hin",
            "is":"ice",
            "id":"ind",
            "it":"ita",
            "ja":"jpn",
            "kk":"kaz",
            "ko":"kor",
            "lv":"lav",
            "lt":"lit",
            "lb":"ltz",
            "mk":"mac",
            "ms":"may",
            "no":"nor",
            "oc":"oci",
            "fa":"per",
            "ro":"rum",
            "ru":"rus",
            "sr":"scc",
            "sl":"slv",
            "sv":"swe",
            "th":"tha",
            "tr":"tur",
            "uk":"ukr",
            "vi":"vie"}

class OpenSubtitles(SubtitleDatabase.SubtitleDB):
    url = "http://www.opensubtitles.org/"
    site_name = "OpenSubtitles"

    def __init__(self, config, cache_folder_path):
        super(OpenSubtitles, self).__init__(OS_LANGS)
        self.server_url = 'http://api.opensubtitles.org/xml-rpc'
        self.revertlangs = dict(map(lambda item: (item[1],item[0]), self.langs.items()))

    def process(self, filepath, langs):
        ''' main method to call on the plugin, pass the filename and the wished
        languages and it will query OpenSubtitles.org '''

	subs = []
        filehash = self.hashFile(filepath)
        # disabled this part because getFileSize gives negative values for large files
        # on BoxeeBox, which makes this useless for larger movies
	fname = self.getFileName(filepath)
        if xbmc.getFileSize(filepath):
            log.debug(filehash)
            size = long(xbmc.getFileSize(filepath))
            fname = self.getFileName(filepath)
            subs += self.query(moviehash=filehash, langs=langs, bytesize=size, filename=fname)

        subs += self.query(langs=langs, filename=fname)
	return subs

    def createFile(self, subtitle):
        '''pass the URL of the sub and the file it matches, will unzip it
        and return the path to the created file'''
        suburl = subtitle["link"]
        videofilename = subtitle["filename"]
        srtbasefilename = videofilename.rsplit(".", 1)[0]

	content = self.downloadContent(suburl)
        f = gzip.open(srtbasefilename+".srt.gz")
        dump = open(srtbasefilename+".srt", "wb")
        dump.write(f.read())
        dump.close()
        f.close()
        os.remove(srtbasefilename+".srt.gz")
        return srtbasefilename+".srt"

    def downloadFile(self, link, path):
	content = self.downloadContent(link)
	dump = open(path+".gz", "wb")
        dump.write(content)
        dump.close()
        f = gzip.open(path+".gz")
        dump = open(path, "wb")
        dump.write(f.read())
        dump.close()
        f.close()
        os.remove(path+".gz")
        return path

    def query(self, filename, imdbID=None, moviehash=None, bytesize=None, langs=None):
        ''' Makes a query on opensubtitles and returns info about found subtitles.
            Note: if using moviehash, bytesize is required.    '''
        log.debug('query')
        #Prepare the search
        search = {}
        sublinks = []
	gotHash = False
	fileHash = ''
	if(moviehash):
		gotHash = True
		fileHash = moviehash

        if moviehash:
            search['moviehash'] = moviehash
            gotHash = True
        if imdbID: search['imdbid'] = imdbID
        if bytesize: search['moviebytesize'] = str(bytesize)
        if len(search) == 0:
            log.debug("No search term, we'll use the filename")
            # Let's try to guess what to search:
            guessed_data = self.guessFileData(filename)
            search['query'] = guessed_data['name']
            log.debug(search['query'])

        #Login
        self.server = xmlrpclib.Server(self.server_url)
        socket.setdefaulttimeout(10)
        try:
            log_result = self.server.LogIn("","","eng","periscope")
            log.debug(log_result)
            token = log_result["token"]
        except Exception:
            log.error("Open subtitles could not be contacted for login")
            token = None
            socket.setdefaulttimeout(None)
            return []
        if not token:
            log.error("Open subtitles did not return a token after logging in.")
            return []

        # Search
        self.filename = filename #Used to order the results
        sublinks += self.get_results(token, search, gotHash, fileHash)

        # Logout
        try:
            self.server.LogOut(token)
        except:
            log.error("Open subtitles could not be contacted for logout")
        socket.setdefaulttimeout(None)
        return sublinks

    def get_results(self, token, search, gotHash, fileHash):
        log.debug("query: token='%s', search='%s'" % (token, search))
        try:
            if search:
                results = self.server.SearchSubtitles(token, [search])
        except Exception, e:
            log.error("Could not query the server OpenSubtitles")
            log.debug(e)
            return []
        log.debug("Result: %s" %str(results))

        sublinks = []
        if results['data']:
            log.debug(results['data'])
            # OpenSubtitles hash function is not robust ... We'll use the MovieReleaseName to help us select the best candidate
            for r in sorted(results['data'], self.sort_by_moviereleasename):
                # Only added if the MovieReleaseName matches the file
                result = {}
                result["release"] = r['SubFileName']
                result["link"] = r['SubDownloadLink']
                result["page"] = r['SubDownloadLink']
		if(gotHash == True and r['MovieHash'] == fileHash):
			result["hash"] = True
                result["lang"] = self.getLG(r['SubLanguageID'])
                if search.has_key("query") : #We are using the guessed file name, let's remove some results

                    # We're not actually removing, since due to some BoxeeBox issues we're
                    # always searching by filename. And if a user renames their files (which)
                    # works better with Boxee, then these names will probably never match
                    # so we only prioritize the matches, but keep everything else in there as
                    # well...
                    if r["MovieReleaseName"].startswith(self.filename):
                        sublinks.insert(0,result)
                    else:
                        sublinks.append(result)
                        log.debug("Removing %s because release '%s' has not right start %s" %(result["release"], r["MovieReleaseName"], self.filename))
                else :
                    sublinks.append(result)
        return sublinks

    def sort_by_moviereleasename(self, x, y):
        ''' sorts based on the movierelease name tag. More matching, returns 1'''
        #TODO add also support for subtitles release
        xmatch = x['MovieReleaseName'] and (x['MovieReleaseName'].find(self.filename)>-1 or self.filename.find(x['MovieReleaseName'])>-1)
        ymatch = y['MovieReleaseName'] and (y['MovieReleaseName'].find(self.filename)>-1 or self.filename.find(y['MovieReleaseName'])>-1)
        #print "analyzing %s and %s = %s and %s" %(x['MovieReleaseName'], y['MovieReleaseName'], xmatch, ymatch)
        if xmatch and ymatch:
            if x['MovieReleaseName'] == self.filename or x['MovieReleaseName'].startswith(self.filename) :
                return -1
            return 0
        if not xmatch and not ymatch:
            return 0
        if xmatch and not ymatch:
            return -1
        if not xmatch and ymatch:
            return 1
        return 0
<commit_msg>Better searching, again<commit_after># -*- coding: utf-8 -*-

#   This file is part of periscope.
#   Copyright (c) 2008-2011 Patrick Dessalle <patrick@dessalle.be>
#
#    periscope is free software; you can redistribute it and/or modify
#    it under the terms of the GNU Lesser General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    periscope is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public License
#    along with periscope; if not, write to the Free Software
#    Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

import os, struct, xmlrpclib, commands, gzip, traceback, logging
import xbmc
import socket # For timeout purposes

import SubtitleDatabase

log = logging.getLogger(__name__)

OS_LANGS ={ "en": "eng",
            "fr" : "fre",
            "hu": "hun",
            "cs": "cze",
            "pl" : "pol",
            "sk" : "slo",
            "pt" : "por",
            "pt-br" : "pob",
            "es" : "spa",
            "el" : "ell",
            "ar":"ara",
            'sq':'alb',
            "hy":"arm",
            "ay":"ass",
            "bs":"bos",
            "bg":"bul",
            "ca":"cat",
            "zh":"chi",
            "hr":"hrv",
            "da":"dan",
            "nl":"dut",
            "eo":"epo",
            "et":"est",
            "fi":"fin",
            "gl":"glg",
            "ka":"geo",
            "de":"ger",
            "he":"heb",
            "hi":"hin",
            "is":"ice",
            "id":"ind",
            "it":"ita",
            "ja":"jpn",
            "kk":"kaz",
            "ko":"kor",
            "lv":"lav",
            "lt":"lit",
            "lb":"ltz",
            "mk":"mac",
            "ms":"may",
            "no":"nor",
            "oc":"oci",
            "fa":"per",
            "ro":"rum",
            "ru":"rus",
            "sr":"scc",
            "sl":"slv",
            "sv":"swe",
            "th":"tha",
            "tr":"tur",
            "uk":"ukr",
            "vi":"vie"}

class OpenSubtitles(SubtitleDatabase.SubtitleDB):
    url = "http://www.opensubtitles.org/"
    site_name = "OpenSubtitles"

    def __init__(self, config, cache_folder_path):
        super(OpenSubtitles, self).__init__(OS_LANGS)
        self.server_url = 'http://api.opensubtitles.org/xml-rpc'
        self.revertlangs = dict(map(lambda item: (item[1],item[0]), self.langs.items()))

    def process(self, filepath, langs):
        ''' main method to call on the plugin, pass the filename and the wished
        languages and it will query OpenSubtitles.org '''

	subs = []
        filehash = self.hashFile(filepath)
        # disabled this part because getFileSize gives negative values for large files
        # on BoxeeBox, which makes this useless for larger movies
	fname = self.getFileName(filepath)
        if xbmc.getFileSize(filepath):
            log.debug(filehash)
            size = long(xbmc.getFileSize(filepath))
            fname = self.getFileName(filepath)
            subs += self.query(moviehash=filehash, langs=langs, bytesize=size, filename=fname)

        subs += self.query(langs=langs, filename=fname)
	return subs

    def createFile(self, subtitle):
        '''pass the URL of the sub and the file it matches, will unzip it
        and return the path to the created file'''
        suburl = subtitle["link"]
        videofilename = subtitle["filename"]
        srtbasefilename = videofilename.rsplit(".", 1)[0]

	content = self.downloadContent(suburl)
        f = gzip.open(srtbasefilename+".srt.gz")
        dump = open(srtbasefilename+".srt", "wb")
        dump.write(f.read())
        dump.close()
        f.close()
        os.remove(srtbasefilename+".srt.gz")
        return srtbasefilename+".srt"

    def downloadFile(self, link, path):
	content = self.downloadContent(link)
	dump = open(path+".gz", "wb")
        dump.write(content)
        dump.close()
        f = gzip.open(path+".gz")
        dump = open(path, "wb")
        dump.write(f.read())
        dump.close()
        f.close()
        os.remove(path+".gz")
        return path

    def query(self, filename, imdbID=None, moviehash=None, bytesize=None, langs=None):
        ''' Makes a query on opensubtitles and returns info about found subtitles.
            Note: if using moviehash, bytesize is required.    '''
        log.debug('query')
        #Prepare the search
        search = {}
        sublinks = []
	gotHash = False
	fileHash = ''
	if(moviehash):
		gotHash = True
		fileHash = moviehash

        if moviehash:
            search['moviehash'] = moviehash
            gotHash = True
        if imdbID: search['imdbid'] = imdbID
        if bytesize: search['moviebytesize'] = str(bytesize)
        if len(search) == 0:
            log.debug("No search term, we'll use the filename")
            # Let's try to guess what to search:
            guessed_data = self.guessFileData(filename)
	    searchstring = guessed_data['name']
	    if guessed_data['type'] == 'tvshow':
	    	if(int(guessed_data['season']) < 10):
                	guessed_data['season'] = "0"+str(guessed_data['season'])
            	if(int(guessed_data['episode']) < 10):
                	guessed_data['episode'] = "0"+str(guessed_data['episode'])

		searchstring = "%s S%sE%s" % (guessed_data['name'], guessed_data['season'], guessed_data['episode'])

            search['query'] = searchstring
            log.debug(search['query'])

        #Login
        self.server = xmlrpclib.Server(self.server_url)
        socket.setdefaulttimeout(10)
        try:
            log_result = self.server.LogIn("","","eng","periscope")
            log.debug(log_result)
            token = log_result["token"]
        except Exception:
            log.error("Open subtitles could not be contacted for login")
            token = None
            socket.setdefaulttimeout(None)
            return []
        if not token:
            log.error("Open subtitles did not return a token after logging in.")
            return []

        # Search
        self.filename = filename #Used to order the results
        sublinks += self.get_results(token, search, gotHash, fileHash)

        # Logout
        try:
            self.server.LogOut(token)
        except:
            log.error("Open subtitles could not be contacted for logout")
        socket.setdefaulttimeout(None)
        return sublinks

    def get_results(self, token, search, gotHash, fileHash):
        log.debug("query: token='%s', search='%s'" % (token, search))
        try:
            if search:
                results = self.server.SearchSubtitles(token, [search])
        except Exception, e:
            log.error("Could not query the server OpenSubtitles")
            log.debug(e)
            return []
        log.debug("Result: %s" %str(results))

        sublinks = []
        if results['data']:
            log.debug(results['data'])
            # OpenSubtitles hash function is not robust ... We'll use the MovieReleaseName to help us select the best candidate
            for r in sorted(results['data'], self.sort_by_moviereleasename):
                # Only added if the MovieReleaseName matches the file
                result = {}
                result["release"] = r['SubFileName']
                result["link"] = r['SubDownloadLink']
                result["page"] = r['SubDownloadLink']
		if(gotHash == True and r['MovieHash'] == fileHash):
			result["hash"] = True
                result["lang"] = self.getLG(r['SubLanguageID'])
                if search.has_key("query") : #We are using the guessed file name, let's remove some results

                    # We're not actually removing, since due to some BoxeeBox issues we're
                    # always searching by filename. And if a user renames their files (which)
                    # works better with Boxee, then these names will probably never match
                    # so we only prioritize the matches, but keep everything else in there as
                    # well...
                    if r["MovieReleaseName"].startswith(self.filename):
                        sublinks.insert(0,result)
                    else:
                        sublinks.append(result)
                        log.debug("Removing %s because release '%s' has not right start %s" %(result["release"], r["MovieReleaseName"], self.filename))
                else :
                    sublinks.append(result)
        return sublinks

    def sort_by_moviereleasename(self, x, y):
        ''' sorts based on the movierelease name tag. More matching, returns 1'''
        #TODO add also support for subtitles release
        xmatch = x['MovieReleaseName'] and (x['MovieReleaseName'].find(self.filename)>-1 or self.filename.find(x['MovieReleaseName'])>-1)
        ymatch = y['MovieReleaseName'] and (y['MovieReleaseName'].find(self.filename)>-1 or self.filename.find(y['MovieReleaseName'])>-1)
        #print "analyzing %s and %s = %s and %s" %(x['MovieReleaseName'], y['MovieReleaseName'], xmatch, ymatch)
        if xmatch and ymatch:
            if x['MovieReleaseName'] == self.filename or x['MovieReleaseName'].startswith(self.filename) :
                return -1
            return 0
        if not xmatch and not ymatch:
            return 0
        if xmatch and not ymatch:
            return -1
        if not xmatch and ymatch:
            return 1
        return 0
<|endoftext|>


<commit_before>from __future__ import division
import logging
import numpy as np
from pysweat.persistence.streams import load_stream
from pysweat.transformation.gps import lat_long_to_x_y
from pysweat.transformation.streams import smooth, derivative, rolling_similarity
from pysweat.transformation.similarities import cosine_similarity, cosine_to_deviation
from pysweat.transformation.windows import subtract_n_minutes

class ActivityFeatures(object):
    def __init__(self, database_driver):
        self.database_driver = database_driver

    def __total_deviation(self, activity_id):
        stream_df = load_stream(self.database_driver, activity_id, 'latlng')
        stream_df = lat_long_to_x_y(stream_df)
        stream_df = smooth(stream_df, smooth_colname='x')
        stream_df = smooth(stream_df, smooth_colname='y')
        stream_df = derivative(stream_df, derivative_colname='x_smooth')
        stream_df = derivative(stream_df, derivative_colname='y_smooth')
        stream_df = rolling_similarity(stream_df, cosine_similarity, 'dx_smooth_dt', 'dy_smooth_dt')

        try:
            return np.nansum([cosine_to_deviation(cos)
                              for cos in stream_df.cosine_similarity_dx_smooth_dt_dy_smooth_dt])
        except ValueError:
            logging.warning('Failed to compute route deviation for activity %d, returning NaN' % activity_id)
            return np.nan

    def turns_per_km(self, activity_df):
        activity_df['turns_per_km'] = [self.__total_deviation(activity_df.strava_id[i]) / activity_df.distance[i]
                                       for i in activity_df.index]
        return activity_df

    def max_value_maintained_for_n_minutes(self, activity_df, measurement='heartrate', window_size=5):
        all_max_values = [float('NaN')] * len(activity_df)

        for i in range(len(activity_df)):
            stream_df = load_stream(self.database_driver, activity_df.strava_id[i], '%s' % measurement)
            if stream_df is not None:
                min_index = stream_df.index.min()
                all_max_values[i] = np.nanmax(
                    [stream_df[measurement].loc[subtract_n_minutes(second,
                                                                   minutes=window_size,
                                                                   minimum_value=min_index):second].min()
                     for second in stream_df.index])

        activity_df['max_%s_%d_minutes' % (measurement, window_size)] = all_max_values
        return activity_df
<commit_msg>Added protection against non-existing or empty streams<commit_after>from __future__ import division
import logging
import numpy as np
from pysweat.persistence.streams import load_stream
from pysweat.transformation.gps import lat_long_to_x_y
from pysweat.transformation.streams import smooth, derivative, rolling_similarity
from pysweat.transformation.similarities import cosine_similarity, cosine_to_deviation
from pysweat.transformation.windows import subtract_n_minutes

class ActivityFeatures(object):
    def __init__(self, database_driver):
        self.database_driver = database_driver

    def __total_deviation(self, activity_id):
        stream_df = load_stream(self.database_driver, activity_id, 'latlng')

        if stream_df is not None:
            stream_df = lat_long_to_x_y(stream_df)
            stream_df = smooth(stream_df, smooth_colname='x')
            stream_df = smooth(stream_df, smooth_colname='y')
            stream_df = derivative(stream_df, derivative_colname='x_smooth')
            stream_df = derivative(stream_df, derivative_colname='y_smooth')
            stream_df = rolling_similarity(stream_df, cosine_similarity, 'dx_smooth_dt', 'dy_smooth_dt')

            try:
                return np.nansum([cosine_to_deviation(cos)
                                  for cos in stream_df.cosine_similarity_dx_smooth_dt_dy_smooth_dt])
            except ValueError:
                logging.warning('Failed to compute route deviation for activity %d, returning NaN' % activity_id)
                return np.nan
        else:
            logging.warning('No gps stream available for activity %d, returning NaN' % activity_id)
            return np.nan

    def turns_per_km(self, activity_df):
        activity_df['turns_per_km'] = [self.__total_deviation(activity_df.strava_id[i]) / activity_df.distance[i]
                                       for i in activity_df.index]
        return activity_df

    def max_value_maintained_for_n_minutes(self, activity_df, measurement='heartrate', window_size=5):
        all_max_values = [float('NaN')] * len(activity_df)

        for i in range(len(activity_df)):
            stream_df = load_stream(self.database_driver, activity_df.strava_id[i], '%s' % measurement)
            if stream_df is not None:
                min_index = stream_df.index.min()
                all_max_values[i] = np.nanmax(
                    [stream_df[measurement].loc[subtract_n_minutes(second,
                                                                   minutes=window_size,
                                                                   minimum_value=min_index):second].min()
                     for second in stream_df.index])

        activity_df['max_%s_%d_minutes' % (measurement, window_size)] = all_max_values
        return activity_df
<|endoftext|>


<commit_before>"""
Test suite for openjpeg's callback functions.
"""
# Standard library imports ...
from io import StringIO
import os
import tempfile
import warnings
import unittest
from unittest.mock import patch

# Local imports ...
import glymur

@unittest.skipIf(glymur.version.openjpeg_version[0] != '2',
                 "Missing openjp2 library.")
class TestCallbacks(unittest.TestCase):
    """Test suite for callbacks."""

    def setUp(self):
        self.jp2file = glymur.data.nemo()
        self.j2kfile = glymur.data.goodstuff()

    @unittest.skipIf(os.name == "nt", "Temporary file issue on window.")
    def test_info_callback_on_write_backwards_compatibility(self):
        """Verify messages printed when writing an image in verbose mode."""
        j = glymur.Jp2k(self.jp2file)
        with warnings.catch_warnings():
            # Ignore a library warning.
            warnings.simplefilter('ignore')
            tiledata = j.read(tile=0)
        with tempfile.NamedTemporaryFile(suffix='.jp2') as tfile:
            with patch('sys.stdout', new=StringIO()) as fake_out:
                glymur.Jp2k(tfile.name, data=tiledata, verbose=True)
            actual = fake_out.getvalue().strip()
        expected = '[INFO] tile number 1 / 1'
        self.assertEqual(actual, expected)

    @unittest.skipIf(os.name == "nt", "Temporary file issue on window.")
    def test_info_callback_on_write(self):
        """Verify messages printed when writing an image in verbose mode."""
        j = glymur.Jp2k(self.jp2file)
        tiledata = j[:]
        with tempfile.NamedTemporaryFile(suffix='.jp2') as tfile:
            with patch('sys.stdout', new=StringIO()) as fake_out:
                glymur.Jp2k(tfile.name, data=tiledata, verbose=True)
                actual = fake_out.getvalue().strip()
        expected = '[INFO] tile number 1 / 1'
        self.assertEqual(actual, expected)

    def test_info_callbacks_on_read(self):
        """stdio output when info callback handler is enabled"""

        # Verify that we get the expected stdio output when our internal info
        # callback handler is enabled.
        jp2 = glymur.Jp2k(self.j2kfile)
        with patch('sys.stdout', new=StringIO()) as fake_out:
            jp2.verbose = True
            jp2[::2, ::2]
            actual = fake_out.getvalue().strip()

        expected = ('[INFO] Start to read j2k main header (0).\n'
                    '[INFO] Main header has been correctly decoded.\n'
                    '[INFO] Setting decoding area to 0,0,480,800\n'
                    '[INFO] Header of tile 1 / 1 has been read.\n'
                    '[INFO] Tile 1/1 has been decoded.\n'
                    '[INFO] Image data has been updated with tile 1.')
        self.assertEqual(actual, expected)
<commit_msg>Fix info callback test failure<commit_after>"""
Test suite for openjpeg's callback functions.
"""
# Standard library imports ...
from io import StringIO
import os
import tempfile
import warnings
import unittest
from unittest.mock import patch

# Local imports ...
import glymur

@unittest.skipIf(glymur.version.openjpeg_version[0] != '2',
                 "Missing openjp2 library.")
class TestCallbacks(unittest.TestCase):
    """Test suite for callbacks."""

    def setUp(self):
        self.jp2file = glymur.data.nemo()
        self.j2kfile = glymur.data.goodstuff()

    @unittest.skipIf(os.name == "nt", "Temporary file issue on window.")
    def test_info_callback_on_write_backwards_compatibility(self):
        """Verify messages printed when writing an image in verbose mode."""
        j = glymur.Jp2k(self.jp2file)
        with warnings.catch_warnings():
            # Ignore a library warning.
            warnings.simplefilter('ignore')
            tiledata = j.read(tile=0)
        with tempfile.NamedTemporaryFile(suffix='.jp2') as tfile:
            with patch('sys.stdout', new=StringIO()) as fake_out:
                glymur.Jp2k(tfile.name, data=tiledata, verbose=True)
            actual = fake_out.getvalue().strip()
        expected = '[INFO] tile number 1 / 1'
        self.assertEqual(actual, expected)

    @unittest.skipIf(os.name == "nt", "Temporary file issue on window.")
    def test_info_callback_on_write(self):
        """Verify messages printed when writing an image in verbose mode."""
        j = glymur.Jp2k(self.jp2file)
        tiledata = j[:]
        with tempfile.NamedTemporaryFile(suffix='.jp2') as tfile:
            with patch('sys.stdout', new=StringIO()) as fake_out:
                glymur.Jp2k(tfile.name, data=tiledata, verbose=True)
                actual = fake_out.getvalue().strip()
        expected = '[INFO] tile number 1 / 1'
        self.assertEqual(actual, expected)

    def test_info_callbacks_on_read(self):
        """
        SCENARIO:  the verbose attribute is set to True

        EXPECTED RESULT:  The info callback handler should be enabled.  There
        should be [INFO] output present in sys.stdout.
        """
        jp2 = glymur.Jp2k(self.j2kfile)
        with patch('sys.stdout', new=StringIO()) as fake_out:
            jp2.verbose = True
            jp2[::2, ::2]
            actual = fake_out.getvalue().strip()

        self.assertIn('[INFO]', actual)
<|endoftext|>


<commit_before>__version__ = '4.0.21'
<commit_msg>start developing version 4.0.22<commit_after>__version__ = '4.0.22'
<|endoftext|>