<commit_before>#!/usr/bin/env python
# Copyright (c) 2011 Citrix Systems, Inc.
#
# This program 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; version 2.1 only. with the special
# exception on linking described in file LICENSE.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.

import os
import os.path
import re
import tempfile

COUNTER = 0

class MenuEntry(object):
    def __init__(self, hypervisor, hypervisor_args, kernel, kernel_args,
                 initrd, title = None):
        self.hypervisor = hypervisor
        self.hypervisor_args = hypervisor_args
        self.kernel = kernel
        self.kernel_args = kernel_args
        self.initrd = initrd
        self.title = title

    def getHypervisorArgs(self):
        return re.findall(r'\S[^ "]*(?:"[^"]*")?\S*', self.hypervisor_args)

    def setHypervisorArgs(self, args):
        self.hypervisor_args = ' '.join(args)

    def getKernelArgs(self):
        return re.findall(r'\S[^ "]*(?:"[^"]*")?\S*', self.kernel_args)

    def setKernelArgs(self, args):
        self.kernel_args = ' '.join(args)

class Bootloader(object):
    def __init__(self, src_fmt, src_file, menu = None, menu_order = None,
                 default = None, timeout = None, serial = None,
                 location = None):

        if menu is None:
            menu = {}

        if menu_order is None:
            menu_order = []

        self.src_fmt = src_fmt
        self.src_file = src_file
        self.menu = menu
        self.menu_order = menu_order
        self.default = default
        self.timeout = timeout
        self.serial = serial
        self.location = location and location or 'mbr'

    def append(self, label, entry):
        self.menu[label] = entry
        self.menu_order.append(label)

    def remove(self, label):
        del self.menu[label]
        self.menu_order.remove(label)

    @classmethod
    def readExtLinux(cls, src_file):
        menu = {}
        menu_order = []
        default = None
        timeout = None
        location = None
        serial = None
        label = None
        title = None

        fh = open(src_file)
        try:
            for line in fh:
                l = line.strip()
                els = l.split(None, 2)
                if len(els) == 0:
                    continue

                # header
                if (l.startswith('# location ') and len(els) == 3 and
                    els[2] in ['mbr', 'partition']):
                    location = els[2]
                elif els[0] == 'serial' and len(els) > 1:
                    baud = 9600
                    if len(els) > 2:
                        baud = int(els[2])
                    serial = {'port': int(els[1]), 'baud': baud}
                elif els[0] == 'default' and len(els) == 2:
                    default = els[1]
                elif els[0] == 'timeout' and len(els) == 2:
                    timeout = int(els[1])

                # menu
                elif els[0] == 'label' and len(els) == 2:
                    label = els[1]
                    menu[label] = {}
                    menu_order.append(label)
                    title = None
                elif label:
                    if els[0] == '#':
                        title = l[1:].lstrip()
                    elif els[0] == 'append' and len(els) > 1:
                        # els[2] contains hypervisor args, kernel,
                        # kernel args & initrd
                        args = map(lambda x: x.strip(), els[2].split('---'))
                        if len(args) == 3:
                            kernel = args[1].split(None, 1)
                            if len(kernel) == 2:
                                menu[label] = MenuEntry(els[1], args[0],
                                                        kernel[0], kernel[1],
                                                        args[2], title)
        finally:
            fh.close()

        return cls('extlinux', src_file, menu, menu_order, default, timeout,
                   serial, location)

    @classmethod
    def readGrub(cls, src_file):
        menu = {}
        menu_order = []
        default = 0
        timeout = None
        location = None
        serial = None
        label = None
        title = None
        hypervisor = None
        hypervisor_args = None
        kernel = None
        kernel_args = None

        def create_label(title):
            global COUNTER

            # FIXME use branding
            if title == 'XenServer':
                return 'xe'

            if title.endswith('(Serial)'):
                return 'xe-serial'
            if title.endswith('Safe Mode'):
                return 'safe'
            if ' / ' in title:
                if '(Serial,' in title:
                    return 'fallback-serial'
                else:
                    return 'fallback'
            COUNTER += 1
            return "label%d" % COUNTER

        fh = open(src_file)
        try:
            for line in fh:
                l = line.strip()
                els = l.split(None, 2)
                if len(els) == 0:
                    continue

                # header
                if (l.startswith('# location ') and len(els) == 3 and
                    els[2] in ['mbr', 'partition']):
                    location = els[2]
                elif els[0] == 'serial' and len(els) > 1:
                    port = 0
                    baud = 9600
                    for arg in l.split(None, 1)[1].split():
                        if '=' in arg:
                            opt, val = arg.split('=')
                            if opt == '--unit':
                                port = int(val)
                            elif opt == '--speed':
                                baud = int(val)
                    serial = {'port': port, 'baud': baud}
                elif els[0] == 'default' and len(els) == 2:
                    # default is index into menu list, fixup later
                    default = int(els[1])
                elif els[0] == 'timeout' and len(els) == 2:
                    timeout = int(els[1]) * 10

                # menu
                elif els[0] == 'title' and len(els) > 1:
                    title = l.split(None, 1)[1]
                elif title:
                    if els[0] == 'kernel' and len(els) > 2:
                        hypervisor, hypervisor_args = (l.split(None, 1)
                                                       [1].split(None, 1))
                    elif els[0] == 'module' and len(els) > 1:
                        if kernel and hypervisor:
                            # second module == initrd
                            label = create_label(title)
                            menu_order.append(label)
                            menu[label] = MenuEntry(hypervisor, hypervisor_args,
                                                    kernel, kernel_args,
                                                    els[1], title)
                            hypervisor = None
                            kernel = None
                        else:
                            kernel, kernel_args = (l.split(None, 1)
                                                   [1].split(None, 1))
                    elif els[0] == 'initrd' and len(els) > 1:
                        # not multiboot
                        kernel = hypervisor
                        kernel_args = hypervisor_args
                        label = create_label(title)
                        menu_order.append(label)
                        menu[label] = MenuEntry(None, None,
                                                kernel, kernel_args,
                                                els[1], title)
                        hypervisor = None
                        hypervisor_args = None

            # fixup default
            if len(menu_order) > default:
                default = menu_order[default]
        finally:
            fh.close()

        return cls('grub', src_file, menu, menu_order, default,
                   timeout, serial, location)

    @classmethod
    def loadExisting(cls, root = '/'):
        if os.path.exists(os.path.join(root, "boot/extlinux.conf")):
            return cls.readExtLinux(os.path.join(root, "boot/extlinux.conf"))
        elif os.path.exists(os.path.join(root, "boot/grub/menu.lst")):
            return cls.readGrub(os.path.join(root, "boot/grub/menu.lst"))
        else:
            raise RuntimeError, "No existing bootloader configuration found"

    def writeExtLinux(self, dst_file = None):
        if hasattr(dst_file, 'name'):
            fh = dst_file
        else:
            fh = open(dst_file, 'w')
        print >> fh, "# location " + self.location

        if self.serial:
            print >> fh, "serial %d %d" % (self.serial['port'],
                                           self.serial['baud'])
        if self.default:
            print >> fh, "default " + self.default
        print >> fh, "prompt 1"
        if self.timeout:
            print >> fh, "timeout %d" % self.timeout

        for label in self.menu_order:
            print >> fh, "\nlabel " + label
            m = self.menu[label]
            if m.title:
                print >> fh, "  # " + m.title
            if m.hypervisor:
                print >> fh, "  kernel mboot.c32"
                print >> fh, "  append %s %s --- %s %s --- %s" % \
                    (m.hypervisor, m.hypervisor_args, m.kernel, m.kernel_args, m.initrd)
            else:
                print >> fh, "  kernel " + m.kernel
                print >> fh, "  append " + m.kernel_args
                print >> fh, "  initrd " + m.initrd
        if not hasattr(dst_file, 'name'):
            fh.close()

    def writeGrub(self, dst_file = None):
        if hasattr(dst_file, 'name'):
            fh = dst_file
        else:
            fh = open(dst_file, 'w')
        print >> fh, "# location " + self.location

        if self.serial:
            print >> fh, "serial --unit=%d --speed=%s" % (self.serial['port'],
                                                          self.serial['baud'])
            print >> fh, "terminal --timeout=10 console serial"
        else:
            print >> fh, "terminal console"
        if self.default:
            for i in range(len(self.menu_order)):
                if self.menu_order[i] == self.default:
                    print >> fh, "default %d" % i
                    break
        if self.timeout:
            print >> fh, "timeout %d" % (self.timeout / 10)

        for label in self.menu_order:
            m = self.menu[label]
            print >> fh, "\ntitle " + m.title
            if m.hypervisor:
                print >> fh, "   kernel " + m.hypervisor + " " + m.hypervisor_args
                print >> fh, "   module " + m.kernel + " " + m.kernel_args
                print >> fh, "   module " + m.initrd
            else:
                print >> fh, "   kernel " + m.kernel + " " + m.kernel_args
                print >> fh, "   initrd " + m.initrd
        if not hasattr(dst_file, 'name'):
            fh.close()

    def commit(self, dst_file = None):
        if not dst_file:
            dst_file = self.src_file

        # write to temp file in final destination directory
        fd, tmp_file = tempfile.mkstemp(dir = os.path.dirname(dst_file))

        if self.src_fmt == 'extlinux':
            self.writeExtLinux(tmp_file)
        elif self.src_fmt == 'grub':
            self.writeGrub(tmp_file)

        # atomically replace destination file
        os.close(fd)
        os.rename(tmp_file, dst_file)

    @classmethod
    def newDefault(cls, kernel_link_name, initrd_link_name, root = '/'):
        b = cls.loadExisting(root)
        # FIXME handle initial case
        if b.menu[b.default].kernel != kernel_link_name:
            backup = []
            if not os.path.exists(kernel_link_name):
                raise RuntimeError, "kernel symlink not found"
            if not os.path.exists(initrd_link_name):
                raise RuntimeError, "initrd symlink not found"
            old_kernel_link = b.menu[b.default].kernel
            old_ver = 'old'
            m = re.search(r'(-\d+\.\d+)-', old_kernel_link)
            if m:
                old_ver = m.group(1)

            for k, v in b.menu.items():
                if v.kernel == old_kernel_link:
                    backup.append((k+old_ver, MenuEntry(**v.__dict__)))
                    v.kernel = kernel_link_name
                    v.initrd = initrd_link_name

            if len(backup) > 0:
                for l, e in backup:
                    b.append(l, e)
            b.commit()
<commit_msg>check for links below root<commit_after>#!/usr/bin/env python
# Copyright (c) 2011 Citrix Systems, Inc.
#
# This program 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; version 2.1 only. with the special
# exception on linking described in file LICENSE.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.

import os
import os.path
import re
import tempfile

COUNTER = 0

class MenuEntry(object):
    def __init__(self, hypervisor, hypervisor_args, kernel, kernel_args,
                 initrd, title = None):
        self.hypervisor = hypervisor
        self.hypervisor_args = hypervisor_args
        self.kernel = kernel
        self.kernel_args = kernel_args
        self.initrd = initrd
        self.title = title

    def getHypervisorArgs(self):
        return re.findall(r'\S[^ "]*(?:"[^"]*")?\S*', self.hypervisor_args)

    def setHypervisorArgs(self, args):
        self.hypervisor_args = ' '.join(args)

    def getKernelArgs(self):
        return re.findall(r'\S[^ "]*(?:"[^"]*")?\S*', self.kernel_args)

    def setKernelArgs(self, args):
        self.kernel_args = ' '.join(args)

class Bootloader(object):
    def __init__(self, src_fmt, src_file, menu = None, menu_order = None,
                 default = None, timeout = None, serial = None,
                 location = None):

        if menu is None:
            menu = {}

        if menu_order is None:
            menu_order = []

        self.src_fmt = src_fmt
        self.src_file = src_file
        self.menu = menu
        self.menu_order = menu_order
        self.default = default
        self.timeout = timeout
        self.serial = serial
        self.location = location and location or 'mbr'

    def append(self, label, entry):
        self.menu[label] = entry
        self.menu_order.append(label)

    def remove(self, label):
        del self.menu[label]
        self.menu_order.remove(label)

    @classmethod
    def readExtLinux(cls, src_file):
        menu = {}
        menu_order = []
        default = None
        timeout = None
        location = None
        serial = None
        label = None
        title = None

        fh = open(src_file)
        try:
            for line in fh:
                l = line.strip()
                els = l.split(None, 2)
                if len(els) == 0:
                    continue

                # header
                if (l.startswith('# location ') and len(els) == 3 and
                    els[2] in ['mbr', 'partition']):
                    location = els[2]
                elif els[0] == 'serial' and len(els) > 1:
                    baud = 9600
                    if len(els) > 2:
                        baud = int(els[2])
                    serial = {'port': int(els[1]), 'baud': baud}
                elif els[0] == 'default' and len(els) == 2:
                    default = els[1]
                elif els[0] == 'timeout' and len(els) == 2:
                    timeout = int(els[1])

                # menu
                elif els[0] == 'label' and len(els) == 2:
                    label = els[1]
                    menu[label] = {}
                    menu_order.append(label)
                    title = None
                elif label:
                    if els[0] == '#':
                        title = l[1:].lstrip()
                    elif els[0] == 'append' and len(els) > 1:
                        # els[2] contains hypervisor args, kernel,
                        # kernel args & initrd
                        args = map(lambda x: x.strip(), els[2].split('---'))
                        if len(args) == 3:
                            kernel = args[1].split(None, 1)
                            if len(kernel) == 2:
                                menu[label] = MenuEntry(els[1], args[0],
                                                        kernel[0], kernel[1],
                                                        args[2], title)
        finally:
            fh.close()

        return cls('extlinux', src_file, menu, menu_order, default, timeout,
                   serial, location)

    @classmethod
    def readGrub(cls, src_file):
        menu = {}
        menu_order = []
        default = 0
        timeout = None
        location = None
        serial = None
        label = None
        title = None
        hypervisor = None
        hypervisor_args = None
        kernel = None
        kernel_args = None

        def create_label(title):
            global COUNTER

            # FIXME use branding
            if title == 'XenServer':
                return 'xe'

            if title.endswith('(Serial)'):
                return 'xe-serial'
            if title.endswith('Safe Mode'):
                return 'safe'
            if ' / ' in title:
                if '(Serial,' in title:
                    return 'fallback-serial'
                else:
                    return 'fallback'
            COUNTER += 1
            return "label%d" % COUNTER

        fh = open(src_file)
        try:
            for line in fh:
                l = line.strip()
                els = l.split(None, 2)
                if len(els) == 0:
                    continue

                # header
                if (l.startswith('# location ') and len(els) == 3 and
                    els[2] in ['mbr', 'partition']):
                    location = els[2]
                elif els[0] == 'serial' and len(els) > 1:
                    port = 0
                    baud = 9600
                    for arg in l.split(None, 1)[1].split():
                        if '=' in arg:
                            opt, val = arg.split('=')
                            if opt == '--unit':
                                port = int(val)
                            elif opt == '--speed':
                                baud = int(val)
                    serial = {'port': port, 'baud': baud}
                elif els[0] == 'default' and len(els) == 2:
                    # default is index into menu list, fixup later
                    default = int(els[1])
                elif els[0] == 'timeout' and len(els) == 2:
                    timeout = int(els[1]) * 10

                # menu
                elif els[0] == 'title' and len(els) > 1:
                    title = l.split(None, 1)[1]
                elif title:
                    if els[0] == 'kernel' and len(els) > 2:
                        hypervisor, hypervisor_args = (l.split(None, 1)
                                                       [1].split(None, 1))
                    elif els[0] == 'module' and len(els) > 1:
                        if kernel and hypervisor:
                            # second module == initrd
                            label = create_label(title)
                            menu_order.append(label)
                            menu[label] = MenuEntry(hypervisor, hypervisor_args,
                                                    kernel, kernel_args,
                                                    els[1], title)
                            hypervisor = None
                            kernel = None
                        else:
                            kernel, kernel_args = (l.split(None, 1)
                                                   [1].split(None, 1))
                    elif els[0] == 'initrd' and len(els) > 1:
                        # not multiboot
                        kernel = hypervisor
                        kernel_args = hypervisor_args
                        label = create_label(title)
                        menu_order.append(label)
                        menu[label] = MenuEntry(None, None,
                                                kernel, kernel_args,
                                                els[1], title)
                        hypervisor = None
                        hypervisor_args = None

            # fixup default
            if len(menu_order) > default:
                default = menu_order[default]
        finally:
            fh.close()

        return cls('grub', src_file, menu, menu_order, default,
                   timeout, serial, location)

    @classmethod
    def loadExisting(cls, root = '/'):
        if os.path.exists(os.path.join(root, "boot/extlinux.conf")):
            return cls.readExtLinux(os.path.join(root, "boot/extlinux.conf"))
        elif os.path.exists(os.path.join(root, "boot/grub/menu.lst")):
            return cls.readGrub(os.path.join(root, "boot/grub/menu.lst"))
        else:
            raise RuntimeError, "No existing bootloader configuration found"

    def writeExtLinux(self, dst_file = None):
        if hasattr(dst_file, 'name'):
            fh = dst_file
        else:
            fh = open(dst_file, 'w')
        print >> fh, "# location " + self.location

        if self.serial:
            print >> fh, "serial %d %d" % (self.serial['port'],
                                           self.serial['baud'])
        if self.default:
            print >> fh, "default " + self.default
        print >> fh, "prompt 1"
        if self.timeout:
            print >> fh, "timeout %d" % self.timeout

        for label in self.menu_order:
            print >> fh, "\nlabel " + label
            m = self.menu[label]
            if m.title:
                print >> fh, "  # " + m.title
            if m.hypervisor:
                print >> fh, "  kernel mboot.c32"
                print >> fh, "  append %s %s --- %s %s --- %s" % \
                    (m.hypervisor, m.hypervisor_args, m.kernel, m.kernel_args, m.initrd)
            else:
                print >> fh, "  kernel " + m.kernel
                print >> fh, "  append " + m.kernel_args
                print >> fh, "  initrd " + m.initrd
        if not hasattr(dst_file, 'name'):
            fh.close()

    def writeGrub(self, dst_file = None):
        if hasattr(dst_file, 'name'):
            fh = dst_file
        else:
            fh = open(dst_file, 'w')
        print >> fh, "# location " + self.location

        if self.serial:
            print >> fh, "serial --unit=%d --speed=%s" % (self.serial['port'],
                                                          self.serial['baud'])
            print >> fh, "terminal --timeout=10 console serial"
        else:
            print >> fh, "terminal console"
        if self.default:
            for i in range(len(self.menu_order)):
                if self.menu_order[i] == self.default:
                    print >> fh, "default %d" % i
                    break
        if self.timeout:
            print >> fh, "timeout %d" % (self.timeout / 10)

        for label in self.menu_order:
            m = self.menu[label]
            print >> fh, "\ntitle " + m.title
            if m.hypervisor:
                print >> fh, "   kernel " + m.hypervisor + " " + m.hypervisor_args
                print >> fh, "   module " + m.kernel + " " + m.kernel_args
                print >> fh, "   module " + m.initrd
            else:
                print >> fh, "   kernel " + m.kernel + " " + m.kernel_args
                print >> fh, "   initrd " + m.initrd
        if not hasattr(dst_file, 'name'):
            fh.close()

    def commit(self, dst_file = None):
        if not dst_file:
            dst_file = self.src_file

        # write to temp file in final destination directory
        fd, tmp_file = tempfile.mkstemp(dir = os.path.dirname(dst_file))

        if self.src_fmt == 'extlinux':
            self.writeExtLinux(tmp_file)
        elif self.src_fmt == 'grub':
            self.writeGrub(tmp_file)

        # atomically replace destination file
        os.close(fd)
        os.rename(tmp_file, dst_file)

    @classmethod
    def newDefault(cls, kernel_link_name, initrd_link_name, root = '/'):
        b = cls.loadExisting(root)
        # FIXME handle initial case
        if b.menu[b.default].kernel != kernel_link_name:
            backup = []
            if not os.path.exists(os.path.join(root, kernel_link_name[1:])):
                raise RuntimeError, "kernel symlink not found"
            if not os.path.exists(os.path.join(root, initrd_link_name[1:])):
                raise RuntimeError, "initrd symlink not found"
            old_kernel_link = b.menu[b.default].kernel
            old_ver = 'old'
            m = re.search(r'(-\d+\.\d+)-', old_kernel_link)
            if m:
                old_ver = m.group(1)

            for k, v in b.menu.items():
                if v.kernel == old_kernel_link:
                    backup.append((k+old_ver, MenuEntry(**v.__dict__)))
                    v.kernel = kernel_link_name
                    v.initrd = initrd_link_name

            if len(backup) > 0:
                for l, e in backup:
                    b.append(l, e)
            b.commit()
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
This ao2mo module is kept for backward compatiblity.  It's recommended to use
pyscf.pbc.df module to get 2e MO integrals
'''
import numpy as np
from pyscf.pbc import df

def general(cell, mo_coeffs, kpts=None, compact=False):
    '''pyscf-style wrapper to get MO 2-el integrals.'''
    if kpts is not None:
        assert len(kpts) == 4
    return df.FFTDF(cell).ao2mo(mo_coeffs, kpts, compact)

def get_mo_eri(cell, mo_coeffs, kpts=None):
    '''Convenience function to return MO 2-el integrals.'''
    return general(cell, mo_coeffs, kpts)

def get_mo_pairs_G(cell, mo_coeffs, kpts=None):
    '''Calculate forward (G|ij) FFT of all MO pairs.

    TODO: - Implement simplifications for real orbitals.

    Args:
        mo_coeff: length-2 list of (nao,nmo) ndarrays
            The two sets of MO coefficients to use in calculating the
            product |ij).

    Returns:
        mo_pairs_G : (ngs, nmoi*nmoj) ndarray
            The FFT of the real-space MO pairs.
    '''
    return df.FFTDF(cell).get_mo_pairs(mo_coeffs, kpts)

def get_mo_pairs_invG(cell, mo_coeffs, kpts=None):
    '''Calculate "inverse" (ij|G) FFT of all MO pairs.

    TODO: - Implement simplifications for real orbitals.

    Args:
        mo_coeff: length-2 list of (nao,nmo) ndarrays
            The two sets of MO coefficients to use in calculating the
            product |ij).

    Returns:
        mo_pairs_invG : (ngs, nmoi*nmoj) ndarray
            The inverse FFTs of the real-space MO pairs.
    '''
    if kpts is None: kpts = numpy.zeros((2,3))
    mo_pairs_G = df.FFTDF(cell).get_mo_pairs((mo_coeffs[1],mo_coeffs[0]),
                                             (kpts[1],kpts[0]))
    nmo0 = mo_coeffs[0].shape[1]
    nmo1 = mo_coeffs[1].shape[1]
    mo_pairs_invG = mo_pairs_G.T.reshape(nmo1,nmo0,-1).transpose(1,0,2).conj()
    mo_pairs_invG = mo_pairs_invG.reshape(nmo0*nmo1,-1).T
    return mo_pairs_invG

def get_ao_pairs_G(cell, kpts=None):
    '''Calculate forward (G|ij) and "inverse" (ij|G) FFT of all AO pairs.

    Args:
        cell : instance of :class:`Cell`

    Returns:
        ao_pairs_G, ao_pairs_invG : (ngs, nao*(nao+1)/2) ndarray
            The FFTs of the real-space AO pairs.

    '''
    return df.FFTDF(cell).get_ao_pairs(kpts)

def get_ao_eri(cell, kpts=None):
    '''Convenience function to return AO 2-el integrals.'''
    if kpts is not None:
        assert len(kpts) == 4
    return df.FFTDF(cell).get_eri(kpts)

if __name__ == '__main__':
    from pyscf.pbc import gto as pgto
    cell = pgto.Cell()
    cell.atom = 'He 1. .5 .5; He .1 1.3 2.1'
    cell.basis = 'ccpvdz'
    cell.a = np.eye(3) * 4.
    cell.gs = [5,5,5]
    cell.build()
    print(get_ao_eri(cell).shape)
<commit_msg>reverted to old pbc eris.py until DF routine is fixed<commit_after>#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
This ao2mo module is kept for backward compatiblity.  It's recommended to use
pyscf.pbc.df module to get 2e MO integrals
'''
import numpy as np
from pyscf.pbc import df
from pyscf.pbc import lib as pbclib
from pyscf.pbc.dft.gen_grid import gen_uniform_grids
from pyscf.pbc.dft.numint import eval_ao
from pyscf.pbc import tools
from pyscf.lib import logger

einsum = pbclib.einsum

def general(cell, mo_coeffs, kpts=None, compact=False):
    '''pyscf-style wrapper to get MO 2-el integrals.'''
    if kpts is not None:
        assert len(kpts) == 4
    return get_mo_eri(cell, mo_coeffs, kpts)
    #return df.FFTDF(cell).ao2mo(mo_coeffs, kpts, compact)

def get_mo_eri(cell, mo_coeffs, kpts=None):
    '''Convenience function to return MO 2-el integrals.'''
    mo_coeff12 = mo_coeffs[:2]
    mo_coeff34 = mo_coeffs[2:]
    if kpts is None:
        kpts12 = kpts34 = q = None
    else:
        kpts12 = kpts[:2]
        kpts34 = kpts[2:]
        q = kpts12[0] - kpts12[1]
        #q = kpts34[1] - kpts34[0]
    if q is None:
        q = np.zeros(3)

    mo_pairs12_kG = get_mo_pairs_G(cell, mo_coeff12, kpts12)
    mo_pairs34_invkG = get_mo_pairs_invG(cell, mo_coeff34, kpts34, q)
    return assemble_eri(cell, mo_pairs12_kG, mo_pairs34_invkG, q)

def get_mo_pairs_G(cell, mo_coeffs, kpts=None, q=None):
    '''Calculate forward (G|ij) FFT of all MO pairs.
    TODO: - Implement simplifications for real orbitals.
    Args:
        mo_coeff: length-2 list of (nao,nmo) ndarrays
            The two sets of MO coefficients to use in calculating the
            product |ij).
    Returns:
        mo_pairs_G : (ngs, nmoi*nmoj) ndarray
            The FFT of the real-space MO pairs.
    '''
    coords = gen_uniform_grids(cell)
    if kpts is None:
        q = np.zeros(3)
        aoR = eval_ao(cell, coords)
        ngs = aoR.shape[0]

        if np.array_equal(mo_coeffs[0], mo_coeffs[1]):
            nmoi = nmoj = mo_coeffs[0].shape[1]
            moiR = mojR = einsum('ri,ia->ra', aoR, mo_coeffs[0])
        else:
            nmoi = mo_coeffs[0].shape[1]
            nmoj = mo_coeffs[1].shape[1]
            moiR = einsum('ri,ia->ra', aoR, mo_coeffs[0])
            mojR = einsum('ri,ia->ra', aoR, mo_coeffs[1])

    else:
        if q is None:
            q = kpts[1]-kpts[0]
        aoR_ki = eval_ao(cell, coords, kpt=kpts[0])
        aoR_kj = eval_ao(cell, coords, kpt=kpts[1])
        ngs = aoR_ki.shape[0]

        nmoi = mo_coeffs[0].shape[1]
        nmoj = mo_coeffs[1].shape[1]
        moiR = einsum('ri,ia->ra', aoR_ki, mo_coeffs[0])
        mojR = einsum('ri,ia->ra', aoR_kj, mo_coeffs[1])

    #mo_pairs_R = einsum('ri,rj->rij', np.conj(moiR), mojR)
    mo_pairs_G = np.zeros([ngs,nmoi*nmoj], np.complex128)

    expmikr = np.exp(-1j*np.dot(q,coords.T))
    for i in xrange(nmoi):
        for j in xrange(nmoj):
            mo_pairs_R_ij = np.conj(moiR[:,i])*mojR[:,j]
            mo_pairs_G[:,i*nmoj+j] = tools.fftk(mo_pairs_R_ij, cell.gs,
                                                expmikr)

    return mo_pairs_G

def get_mo_pairs_invG(cell, mo_coeffs, kpts=None, q=None):
    '''Calculate "inverse" (ij|G) FFT of all MO pairs.
    TODO: - Implement simplifications for real orbitals.
    Args:
        mo_coeff: length-2 list of (nao,nmo) ndarrays
            The two sets of MO coefficients to use in calculating the
            product |ij).
    Returns:
        mo_pairs_invG : (ngs, nmoi*nmoj) ndarray
            The inverse FFTs of the real-space MO pairs.
    '''
    coords = gen_uniform_grids(cell)
    if kpts is None:
        q = np.zeros(3)
        aoR = eval_ao(cell, coords)
        ngs = aoR.shape[0]

        if np.array_equal(mo_coeffs[0], mo_coeffs[1]):
            nmoi = nmoj = mo_coeffs[0].shape[1]
            moiR = mojR = einsum('ri,ia->ra', aoR, mo_coeffs[0])
        else:
            nmoi = mo_coeffs[0].shape[1]
            nmoj = mo_coeffs[1].shape[1]
            moiR = einsum('ri,ia->ra', aoR, mo_coeffs[0])
            mojR = einsum('ri,ia->ra', aoR, mo_coeffs[1])

    else:
        if q is None:
            q = kpts[1]-kpts[0]
        aoR_ki = eval_ao(cell, coords, kpt=kpts[0])
        aoR_kj = eval_ao(cell, coords, kpt=kpts[1])
        ngs = aoR_ki.shape[0]

        nmoi = mo_coeffs[0].shape[1]
        nmoj = mo_coeffs[1].shape[1]
        moiR = einsum('ri,ia->ra', aoR_ki, mo_coeffs[0])
        mojR = einsum('ri,ia->ra', aoR_kj, mo_coeffs[1])

    #mo_pairs_R = einsum('ri,rj->rij', np.conj(moiR), mojR)
    mo_pairs_invG = np.zeros([ngs,nmoi*nmoj], np.complex128)

    expmikr = np.exp(-1j*np.dot(q,coords.T))
    for i in xrange(nmoi):
        for j in xrange(nmoj):
            mo_pairs_R_ij = np.conj(moiR[:,i])*mojR[:,j]
            mo_pairs_invG[:,i*nmoj+j] = np.conj(tools.fftk(np.conj(mo_pairs_R_ij), cell.gs,
                                                           expmikr.conj()))

    return mo_pairs_invG

def assemble_eri(cell, orb_pair_invG1, orb_pair_G2, q=None, verbose=logger.INFO):
    '''Assemble 4-index electron repulsion integrals.
    Returns:
        (nmo1*nmo2, nmo3*nmo4) ndarray
    '''
    log = logger.Logger
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(cell.stdout, verbose)

    log.debug('Performing periodic ERI assembly of (%i, %i) ij,kl pairs',
              orb_pair_invG1.shape[1], orb_pair_G2.shape[1])
    if q is None:
        q = np.zeros(3)

    coulqG = tools.get_coulG(cell, -1.0*q)
    ngs = orb_pair_invG1.shape[0]
    Jorb_pair_G2 = np.einsum('g,gn->gn',coulqG,orb_pair_G2)*(cell.vol/ngs**2)
    eri = np.dot(orb_pair_invG1.T, Jorb_pair_G2)
    return eri

#def get_mo_eri(cell, mo_coeffs, kpts=None):
#    '''Convenience function to return MO 2-el integrals.'''
#    return general(cell, mo_coeffs, kpts)

#def get_mo_pairs_G(cell, mo_coeffs, kpts=None):
#    '''Calculate forward (G|ij) FFT of all MO pairs.
#
#    TODO: - Implement simplifications for real orbitals.
#
#    Args:
#        mo_coeff: length-2 list of (nao,nmo) ndarrays
#            The two sets of MO coefficients to use in calculating the
#            product |ij).
#
#    Returns:
#        mo_pairs_G : (ngs, nmoi*nmoj) ndarray
#            The FFT of the real-space MO pairs.
#    '''
#    return df.FFTDF(cell).get_mo_pairs(mo_coeffs, kpts)
#
#def get_mo_pairs_invG(cell, mo_coeffs, kpts=None):
#    '''Calculate "inverse" (ij|G) FFT of all MO pairs.
#
#    TODO: - Implement simplifications for real orbitals.
#
#    Args:
#        mo_coeff: length-2 list of (nao,nmo) ndarrays
#            The two sets of MO coefficients to use in calculating the
#            product |ij).
#
#    Returns:
#        mo_pairs_invG : (ngs, nmoi*nmoj) ndarray
#            The inverse FFTs of the real-space MO pairs.
#    '''
#    if kpts is None: kpts = numpy.zeros((2,3))
#    mo_pairs_G = df.FFTDF(cell).get_mo_pairs((mo_coeffs[1],mo_coeffs[0]),
#                                             (kpts[1],kpts[0]))
#    nmo0 = mo_coeffs[0].shape[1]
#    nmo1 = mo_coeffs[1].shape[1]
#    mo_pairs_invG = mo_pairs_G.T.reshape(nmo1,nmo0,-1).transpose(1,0,2).conj()
#    mo_pairs_invG = mo_pairs_invG.reshape(nmo0*nmo1,-1).T
#    return mo_pairs_invG

def get_ao_pairs_G(cell, kpts=None):
    '''Calculate forward (G|ij) and "inverse" (ij|G) FFT of all AO pairs.

    Args:
        cell : instance of :class:`Cell`

    Returns:
        ao_pairs_G, ao_pairs_invG : (ngs, nao*(nao+1)/2) ndarray
            The FFTs of the real-space AO pairs.

    '''
    return df.FFTDF(cell).get_ao_pairs(kpts)

def get_ao_eri(cell, kpts=None):
    '''Convenience function to return AO 2-el integrals.'''
    if kpts is not None:
        assert len(kpts) == 4
    return df.FFTDF(cell).get_eri(kpts)

if __name__ == '__main__':
    from pyscf.pbc import gto as pgto
    cell = pgto.Cell()
    cell.atom = 'He 1. .5 .5; He .1 1.3 2.1'
    cell.basis = 'ccpvdz'
    cell.a = np.eye(3) * 4.
    cell.gs = [5,5,5]
    cell.build()
    print(get_ao_eri(cell).shape)
<|endoftext|>


<commit_before>from datetime import datetime

import django_filters
from django.apps import apps
from django.conf import settings
from django.contrib import messages
from django.contrib.auth import get_user_model
from django.db.models import Q
from django.shortcuts import get_object_or_404
from django.shortcuts import redirect
from django.utils.translation import ugettext_lazy as _
from django.utils.translation import ungettext
from django.views import generic
from rules.contrib.views import LoginRequiredMixin

from adhocracy4.filters import views as filter_views
from adhocracy4.filters import widgets as filters_widgets
from adhocracy4.filters.filters import DefaultsFilterSet
from adhocracy4.filters.filters import FreeTextFilter
from adhocracy4.filters.widgets import DropdownLinkWidget
from adhocracy4.projects import models as project_models
from meinberlin.apps.dashboard2 import mixins as a4dashboard_mixins

from . import forms
from . import models

User = get_user_model()

class OrderingWidget(DropdownLinkWidget):
    label = _('Ordering')
    right = True

class OrganisationWidget(DropdownLinkWidget):
    label = _('Organisation')

class FreeTextFilterWidget(filters_widgets.FreeTextFilterWidget):
    label = _('Search')

class ArchivedWidget(DropdownLinkWidget):
    label = _('Archived')

    def __init__(self, attrs=None):
        choices = (
            ('', _('All')),
            ('false', _('No')),
            ('true', _('Yes')),
        )
        super().__init__(attrs, choices)

class YearWidget(DropdownLinkWidget):
    label = _('Year')

    def __init__(self, attrs=None):
        choices = (('', _('Any')),)
        now = datetime.now().year
        try:
            first_year = project_models.Project.objects.earliest('created').\
                created.year
        except project_models.Project.DoesNotExist:
            first_year = now
        for year in range(now, first_year - 1, -1):
            choices += (year, year),
        super().__init__(attrs, choices)

class ProjectFilterSet(DefaultsFilterSet):

    defaults = {
        'is_archived': 'false'
    }

    ordering = django_filters.OrderingFilter(
        choices=(
            ('-created', _('Most recent')),
        ),
        empty_label=None,
        widget=OrderingWidget,
    )

    search = FreeTextFilter(
        widget=FreeTextFilterWidget,
        fields=['name', 'description']
    )

    organisation = django_filters.ModelChoiceFilter(
        queryset=apps.get_model(settings.A4_ORGANISATIONS_MODEL).objects
                     .order_by('name'),
        widget=OrganisationWidget,
    )

    is_archived = django_filters.BooleanFilter(
        widget=ArchivedWidget
    )

    created = django_filters.NumberFilter(
        name='created',
        lookup_expr='year',
        widget=YearWidget,
    )

    class Meta:
        model = project_models.Project
        fields = ['search', 'organisation', 'is_archived', 'created']

class ProjectListView(filter_views.FilteredListView):
    model = project_models.Project
    paginate_by = 16
    filter_set = ProjectFilterSet

    def get_queryset(self):
        return super().get_queryset().filter(
            Q(is_draft=False, is_public=False,
              participants__pk=self.request.user.pk) |
            Q(is_draft=False, is_public=False,
              organisation__initiators__pk=self.request.user.pk) |
            Q(is_draft=False, is_public=True)
        ).distinct()

class ParticipantInviteDetailView(generic.DetailView):
    model = models.ParticipantInvite
    slug_field = 'token'
    slug_url_kwarg = 'invite_token'

    def dispatch(self, request, invite_token, *args, **kwargs):
        if request.user.is_authenticated():
            return redirect(
                'project-participant-invite-update',
                invite_token=invite_token
            )
        else:
            return super().dispatch(request, *args, **kwargs)

class ParticipantInviteUpdateView(LoginRequiredMixin, generic.UpdateView):
    model = models.ParticipantInvite
    form_class = forms.ParticipantInviteForm
    slug_field = 'token'
    slug_url_kwarg = 'invite_token'

    def form_valid(self, form):
        if form.is_accepted():
            form.instance.accept(self.request.user)
            return redirect(form.instance.project.get_absolute_url())
        else:
            form.instance.reject()
            return redirect('/')

class ModeratorInviteDetailView(generic.DetailView):
    model = models.ModeratorInvite
    slug_field = 'token'
    slug_url_kwarg = 'invite_token'

    def dispatch(self, request, invite_token, *args, **kwargs):
        if request.user.is_authenticated():
            return redirect(
                'project-moderator-invite-update',
                invite_token=invite_token
            )
        else:
            return super().dispatch(request, *args, **kwargs)

class ModeratorInviteUpdateView(LoginRequiredMixin, generic.UpdateView):
    model = models.ModeratorInvite
    form_class = forms.ModeratorInviteForm
    slug_field = 'token'
    slug_url_kwarg = 'invite_token'

    def form_valid(self, form):
        if form.is_accepted():
            form.instance.accept(self.request.user)
            return redirect(form.instance.project.get_absolute_url())
        else:
            form.instance.reject()
            return redirect('/')

class AbstractProjectUserInviteListView(
        a4dashboard_mixins.DashboardComponentMixin,
        a4dashboard_mixins.DashboardBaseMixin,
        a4dashboard_mixins.DashboardContextMixin,
        generic.base.TemplateResponseMixin,
        generic.edit.FormMixin,
        generic.detail.SingleObjectMixin,
        generic.edit.ProcessFormView):

    form_class = forms.InviteUsersFromEmailForm
    invite_model = None

    def get(self, request, *args, **kwargs):
        self.object = self.get_object()
        return super().get(request, *args, **kwargs)

    def post(self, request, *args, **kwargs):
        self.object = self.get_object()
        if 'submit_action' in request.POST:
            if request.POST['submit_action'] == 'remove_user':
                pk = int(request.POST['user_pk'])
                user = get_object_or_404(User, pk=pk)
                related_users = getattr(self.object, self.related_users_field)
                related_users.remove(user)
                messages.success(request, self.success_message_removal)
            elif request.POST['submit_action'] == 'remove_invite':
                pk = int(request.POST['invite_pk'])
                invite = self.invite_model.objects.get(pk=pk)
                invite.delete()
                messages.success(request, _('Invitation succesfully removed.'))

            return redirect(self.get_success_url())
        else:
            return super().post(request, *args, **kwargs)

    def filter_existing(self, emails):
        related_users = getattr(self.object, self.related_users_field)
        related_emails = [u.email for u in related_users.all()]
        existing = []
        filtered_emails = []
        for email in emails:
            if email in related_emails:
                existing.append(email)
            else:
                filtered_emails.append(email)
        return filtered_emails, existing

    def filter_pending(self, emails):
        pending = []
        filtered_emails = []
        for email in emails:
            if self.invite_model.objects.filter(email=email).exists():
                pending.append(email)
            else:
                filtered_emails.append(email)
        return filtered_emails, pending

    def form_valid(self, form):
        emails = form.cleaned_data['add_users']

        emails, existing = self.filter_existing(emails)
        if existing:
            messages.error(
                self.request,
                _('Following users already accepted an invitation: ') +
                ', '.join(existing)
            )

        emails, pending = self.filter_pending(emails)
        if pending:
            messages.error(
                self.request,
                _('Following users are already invited: ') +
                ', '.join(pending)
            )

        for email in emails:
            self.invite_model.objects.invite(
                self.request.user,
                self.project,
                email
            )

        messages.success(
            self.request,
            ungettext(self.success_message[0], self.success_message[1],
                      len(emails)).format(len(emails))
        )

        return redirect(self.get_success_url())

    def get_form_kwargs(self):
        kwargs = super().get_form_kwargs()
        kwargs['label'] = self.add_user_field_label
        return kwargs

class DashboardProjectModeratorsView(AbstractProjectUserInviteListView):

    model = project_models.Project
    slug_url_kwarg = 'project_slug'
    template_name = 'meinberlin_projects/project_moderators.html'
    permission_required = 'a4projects.add_project'
    menu_item = 'project'

    related_users_field = 'moderators'
    add_user_field_label = _('Invite moderators via email')
    success_message = ('{} moderator invited.', '{} moderators invited.')
    success_message_removal = _('Moderator successfully removed.')

    invite_model = models.ModeratorInvite

class DashboardProjectParticipantsView(AbstractProjectUserInviteListView):

    model = project_models.Project
    slug_url_kwarg = 'project_slug'
    template_name = 'meinberlin_projects/project_participants.html'
    permission_required = 'a4projects.add_project'
    menu_item = 'project'

    related_users_field = 'participants'
    add_user_field_label = _('Invite users via email')
    success_message = ('{} participant invited.', '{} participants invited.')
    success_message_removal = _('Participant successfully removed.')

    invite_model = models.ParticipantInvite
<commit_msg>refactor project list queryset<commit_after>from datetime import datetime

import django_filters
from django.apps import apps
from django.conf import settings
from django.contrib import messages
from django.contrib.auth import get_user_model
from django.db.models import Q
from django.shortcuts import get_object_or_404
from django.shortcuts import redirect
from django.utils.translation import ugettext_lazy as _
from django.utils.translation import ungettext
from django.views import generic
from rules.contrib.views import LoginRequiredMixin

from adhocracy4.filters import views as filter_views
from adhocracy4.filters import widgets as filters_widgets
from adhocracy4.filters.filters import DefaultsFilterSet
from adhocracy4.filters.filters import FreeTextFilter
from adhocracy4.filters.widgets import DropdownLinkWidget
from adhocracy4.projects import models as project_models
from meinberlin.apps.dashboard2 import mixins as a4dashboard_mixins

from . import forms
from . import models

User = get_user_model()

class OrderingWidget(DropdownLinkWidget):
    label = _('Ordering')
    right = True

class OrganisationWidget(DropdownLinkWidget):
    label = _('Organisation')

class FreeTextFilterWidget(filters_widgets.FreeTextFilterWidget):
    label = _('Search')

class ArchivedWidget(DropdownLinkWidget):
    label = _('Archived')

    def __init__(self, attrs=None):
        choices = (
            ('', _('All')),
            ('false', _('No')),
            ('true', _('Yes')),
        )
        super().__init__(attrs, choices)

class YearWidget(DropdownLinkWidget):
    label = _('Year')

    def __init__(self, attrs=None):
        choices = (('', _('Any')),)
        now = datetime.now().year
        try:
            first_year = project_models.Project.objects.earliest('created').\
                created.year
        except project_models.Project.DoesNotExist:
            first_year = now
        for year in range(now, first_year - 1, -1):
            choices += (year, year),
        super().__init__(attrs, choices)

class ProjectFilterSet(DefaultsFilterSet):

    defaults = {
        'is_archived': 'false'
    }

    ordering = django_filters.OrderingFilter(
        choices=(
            ('-created', _('Most recent')),
        ),
        empty_label=None,
        widget=OrderingWidget,
    )

    search = FreeTextFilter(
        widget=FreeTextFilterWidget,
        fields=['name', 'description']
    )

    organisation = django_filters.ModelChoiceFilter(
        queryset=apps.get_model(settings.A4_ORGANISATIONS_MODEL).objects
                     .order_by('name'),
        widget=OrganisationWidget,
    )

    is_archived = django_filters.BooleanFilter(
        widget=ArchivedWidget
    )

    created = django_filters.NumberFilter(
        name='created',
        lookup_expr='year',
        widget=YearWidget,
    )

    class Meta:
        model = project_models.Project
        fields = ['search', 'organisation', 'is_archived', 'created']

class ProjectListView(filter_views.FilteredListView):
    model = project_models.Project
    paginate_by = 16
    filter_set = ProjectFilterSet

    def get_queryset(self):
        # FIXME: has to be in sync with a4projects.view_project and should
        #        be implemented on the ProjectQueryManager in core

        q_false = Q(pk=None)
        is_superuser = ~q_false if self.request.user.is_superuser else q_false

        return super().get_queryset().filter(
            Q(is_draft=False) & (
                Q(is_public=True) |
                is_superuser |
                Q(participants__pk=self.request.user.pk) |
                Q(organisation__initiators__pk=self.request.user.pk) |
                Q(moderators__pk=self.request.user.pk)
            )
        ).distinct()

class ParticipantInviteDetailView(generic.DetailView):
    model = models.ParticipantInvite
    slug_field = 'token'
    slug_url_kwarg = 'invite_token'

    def dispatch(self, request, invite_token, *args, **kwargs):
        if request.user.is_authenticated():
            return redirect(
                'project-participant-invite-update',
                invite_token=invite_token
            )
        else:
            return super().dispatch(request, *args, **kwargs)

class ParticipantInviteUpdateView(LoginRequiredMixin, generic.UpdateView):
    model = models.ParticipantInvite
    form_class = forms.ParticipantInviteForm
    slug_field = 'token'
    slug_url_kwarg = 'invite_token'

    def form_valid(self, form):
        if form.is_accepted():
            form.instance.accept(self.request.user)
            return redirect(form.instance.project.get_absolute_url())
        else:
            form.instance.reject()
            return redirect('/')

class ModeratorInviteDetailView(generic.DetailView):
    model = models.ModeratorInvite
    slug_field = 'token'
    slug_url_kwarg = 'invite_token'

    def dispatch(self, request, invite_token, *args, **kwargs):
        if request.user.is_authenticated():
            return redirect(
                'project-moderator-invite-update',
                invite_token=invite_token
            )
        else:
            return super().dispatch(request, *args, **kwargs)

class ModeratorInviteUpdateView(LoginRequiredMixin, generic.UpdateView):
    model = models.ModeratorInvite
    form_class = forms.ModeratorInviteForm
    slug_field = 'token'
    slug_url_kwarg = 'invite_token'

    def form_valid(self, form):
        if form.is_accepted():
            form.instance.accept(self.request.user)
            return redirect(form.instance.project.get_absolute_url())
        else:
            form.instance.reject()
            return redirect('/')

class AbstractProjectUserInviteListView(
        a4dashboard_mixins.DashboardComponentMixin,
        a4dashboard_mixins.DashboardBaseMixin,
        a4dashboard_mixins.DashboardContextMixin,
        generic.base.TemplateResponseMixin,
        generic.edit.FormMixin,
        generic.detail.SingleObjectMixin,
        generic.edit.ProcessFormView):

    form_class = forms.InviteUsersFromEmailForm
    invite_model = None

    def get(self, request, *args, **kwargs):
        self.object = self.get_object()
        return super().get(request, *args, **kwargs)

    def post(self, request, *args, **kwargs):
        self.object = self.get_object()
        if 'submit_action' in request.POST:
            if request.POST['submit_action'] == 'remove_user':
                pk = int(request.POST['user_pk'])
                user = get_object_or_404(User, pk=pk)
                related_users = getattr(self.object, self.related_users_field)
                related_users.remove(user)
                messages.success(request, self.success_message_removal)
            elif request.POST['submit_action'] == 'remove_invite':
                pk = int(request.POST['invite_pk'])
                invite = self.invite_model.objects.get(pk=pk)
                invite.delete()
                messages.success(request, _('Invitation succesfully removed.'))

            return redirect(self.get_success_url())
        else:
            return super().post(request, *args, **kwargs)

    def filter_existing(self, emails):
        related_users = getattr(self.object, self.related_users_field)
        related_emails = [u.email for u in related_users.all()]
        existing = []
        filtered_emails = []
        for email in emails:
            if email in related_emails:
                existing.append(email)
            else:
                filtered_emails.append(email)
        return filtered_emails, existing

    def filter_pending(self, emails):
        pending = []
        filtered_emails = []
        for email in emails:
            if self.invite_model.objects.filter(email=email).exists():
                pending.append(email)
            else:
                filtered_emails.append(email)
        return filtered_emails, pending

    def form_valid(self, form):
        emails = form.cleaned_data['add_users']

        emails, existing = self.filter_existing(emails)
        if existing:
            messages.error(
                self.request,
                _('Following users already accepted an invitation: ') +
                ', '.join(existing)
            )

        emails, pending = self.filter_pending(emails)
        if pending:
            messages.error(
                self.request,
                _('Following users are already invited: ') +
                ', '.join(pending)
            )

        for email in emails:
            self.invite_model.objects.invite(
                self.request.user,
                self.project,
                email
            )

        messages.success(
            self.request,
            ungettext(self.success_message[0], self.success_message[1],
                      len(emails)).format(len(emails))
        )

        return redirect(self.get_success_url())

    def get_form_kwargs(self):
        kwargs = super().get_form_kwargs()
        kwargs['label'] = self.add_user_field_label
        return kwargs

class DashboardProjectModeratorsView(AbstractProjectUserInviteListView):

    model = project_models.Project
    slug_url_kwarg = 'project_slug'
    template_name = 'meinberlin_projects/project_moderators.html'
    permission_required = 'a4projects.add_project'
    menu_item = 'project'

    related_users_field = 'moderators'
    add_user_field_label = _('Invite moderators via email')
    success_message = ('{} moderator invited.', '{} moderators invited.')
    success_message_removal = _('Moderator successfully removed.')

    invite_model = models.ModeratorInvite

class DashboardProjectParticipantsView(AbstractProjectUserInviteListView):

    model = project_models.Project
    slug_url_kwarg = 'project_slug'
    template_name = 'meinberlin_projects/project_participants.html'
    permission_required = 'a4projects.add_project'
    menu_item = 'project'

    related_users_field = 'participants'
    add_user_field_label = _('Invite users via email')
    success_message = ('{} participant invited.', '{} participants invited.')
    success_message_removal = _('Participant successfully removed.')

    invite_model = models.ParticipantInvite
<|endoftext|>


<commit_before>import re
import os
import sys
import cffi
import subprocess
import platform

import logging
l = logging.getLogger('cffier')
l.setLevel(logging.DEBUG)

def find_good_bsearch(known_good, questionable):
    ffi = cffi.FFI()
    l.debug("bsearch - trying %d good and %d questionable", len(known_good), len(questionable))

    try:
        ffi.cdef('\n'.join(known_good + questionable))
    except (cffi.CDefError, AssertionError):
        return find_good_bsearch(known_good, questionable[:len(questionable)/2])

    return questionable

def find_good_scan(known_good, questionable):
    ffi = cffi.FFI()
    l.debug("scan - trying %d good and %d questionable", len(known_good), len(questionable))

    #print "GOOD:"
    #print '  ...'
    #print '  ' + '\n  '.join(known_good[-5:])
    #print "UNKNOWN:"
    #print '  ' + '\n  '.join(questionable[:5])
    #print '  ...'
    #print '  ' + '\n  '.join(questionable[-5:])

    try:
        ffi.cdef('\n'.join(known_good + questionable))
        return questionable
    except AssertionError:
        return [ ]
    except cffi.CDefError as e:
        if str(e).count(':') >= 2:
            fail = int(str(e).split('\n')[1].split(':')[1])
        elif 'unrecognized construct' in str(e):
            fail = int(str(e).split()[1][:-1])
        elif 'end of input' in str(e):
            return find_good_scan(known_good, questionable[:-1])
        else:
            raise Exception("Unknown error")
    except cffi.FFIError as e:
        if str(e).count(':') >= 2:
            fail = int(str(e).split('\n')[0].split(':')[1])
        else:
            raise Exception("Unknown error")

    return find_good_scan(known_good, questionable[:fail-2-len(known_good)])

def doit(vex_path):
    cpplist = ['cl', 'cpp']
    cpp = os.getenv("CPP")
    if cpp:
        cpplist.insert(0, cpp)
    if platform.system() == 'Darwin':
        cpplist.insert(0, "clang")

    errs = []
    for cpp in cpplist:
        cmd = [cpp, '-I' + vex_path, os.path.join("pyvex_c", "pyvex.h")]
        if cpp in ('cl', 'clang', 'gcc', 'cc', 'clang++', 'g++'):
            cmd.append("-E")
        try:
            p = subprocess.Popen(cmd,
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE)
            header, stderr = p.communicate()
            if p.returncode != 0 or stderr or not header:
                errs.append((" ".join(cmd), stderr))
                continue
            else:
                break
        except OSError:
            errs.append((" ".join(cmd), -1, "does not exist"))
            continue
    else:
        l.warning("failed commands:\n" +
                  "\n".join("{} -- {}".format(*e) for e in errs))
        raise Exception("Couldn't process pyvex headers." +
                "Please set CPP environmental variable to local path of \"cpp\"." +
                "Note that \"cpp\" and \"g++\" are different."
                )
    #header = vex_pp + pyvex_pp

    linesep = '\r\n' if '\r\n' in header else '\n'
    ffi_text = linesep.join(line for line in header.split(linesep) if '#' not in line and line.strip() != '' and 'jmp_buf' not in line)
    ffi_text = re.sub('\{\s*\} NoOp;', '{ int DONOTUSE; } NoOp;', ffi_text)
    ffi_text = re.sub('__attribute__\s*\(.*\)', '', ffi_text)
    ffi_text = re.sub('__declspec\s*\([^\)]*\)', '', ffi_text)
    ffi_text = ffi_text.replace('__const', 'const')
    ffi_text = ffi_text.replace('__inline', '')
    ffi_text = ffi_text.replace('__w64', '')
    ffi_text = ffi_text.replace('__cdecl', '')
    ffi_text = ffi_text.replace('__int64', 'long')
    ffi_lines = ffi_text.split(linesep)

    good = find_good_scan([], ffi_lines)
    remaining = ffi_lines[len(good)+1:]

    while len(remaining) > 1:
        l.debug("%d uncertain lines remaining", len(remaining))
        new_good = find_good_scan(good, remaining[1:])
        good += new_good
        remaining = remaining[len(new_good)+1:]

    good += [ 'extern VexControl vex_control;' ]

    open('pyvex/vex_ffi.py', 'w').write('ffi_str = """' + '\n'.join(good) + '"""')

if __name__ == '__main__':
    import sys
    logging.basicConfig(level=logging.DEBUG)
    doit(sys.argv[1])
<commit_msg>Remove stderr check from make_ffi.py<commit_after>import re
import os
import sys
import cffi
import subprocess
import platform

import logging
l = logging.getLogger('cffier')
l.setLevel(logging.DEBUG)

def find_good_bsearch(known_good, questionable):
    ffi = cffi.FFI()
    l.debug("bsearch - trying %d good and %d questionable", len(known_good), len(questionable))

    try:
        ffi.cdef('\n'.join(known_good + questionable))
    except (cffi.CDefError, AssertionError):
        return find_good_bsearch(known_good, questionable[:len(questionable)/2])

    return questionable

def find_good_scan(known_good, questionable):
    ffi = cffi.FFI()
    l.debug("scan - trying %d good and %d questionable", len(known_good), len(questionable))

    #print "GOOD:"
    #print '  ...'
    #print '  ' + '\n  '.join(known_good[-5:])
    #print "UNKNOWN:"
    #print '  ' + '\n  '.join(questionable[:5])
    #print '  ...'
    #print '  ' + '\n  '.join(questionable[-5:])

    try:
        ffi.cdef('\n'.join(known_good + questionable))
        return questionable
    except AssertionError:
        return [ ]
    except cffi.CDefError as e:
        if str(e).count(':') >= 2:
            fail = int(str(e).split('\n')[1].split(':')[1])
        elif 'unrecognized construct' in str(e):
            fail = int(str(e).split()[1][:-1])
        elif 'end of input' in str(e):
            return find_good_scan(known_good, questionable[:-1])
        else:
            raise Exception("Unknown error")
    except cffi.FFIError as e:
        if str(e).count(':') >= 2:
            fail = int(str(e).split('\n')[0].split(':')[1])
        else:
            raise Exception("Unknown error")

    return find_good_scan(known_good, questionable[:fail-2-len(known_good)])

def doit(vex_path):
    cpplist = ['cl', 'cpp']
    cpp = os.getenv("CPP")
    if cpp:
        cpplist.insert(0, cpp)
    if platform.system() == 'Darwin':
        cpplist.insert(0, "clang")

    errs = []
    for cpp in cpplist:
        cmd = [cpp, '-I' + vex_path, os.path.join("pyvex_c", "pyvex.h")]
        if cpp in ('cl', 'clang', 'gcc', 'cc', 'clang++', 'g++'):
            cmd.append("-E")
        try:
            p = subprocess.Popen(cmd,
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE)
            header, stderr = p.communicate()
            if p.returncode != 0 or not header:
                errs.append((" ".join(cmd), stderr))
                continue
            else:
                break
        except OSError:
            errs.append((" ".join(cmd), -1, "does not exist"))
            continue
    else:
        l.warning("failed commands:\n" +
                  "\n".join("{} -- {}".format(*e) for e in errs))
        raise Exception("Couldn't process pyvex headers." +
                "Please set CPP environmental variable to local path of \"cpp\"." +
                "Note that \"cpp\" and \"g++\" are different."
                )
    #header = vex_pp + pyvex_pp

    linesep = '\r\n' if '\r\n' in header else '\n'
    ffi_text = linesep.join(line for line in header.split(linesep) if '#' not in line and line.strip() != '' and 'jmp_buf' not in line)
    ffi_text = re.sub('\{\s*\} NoOp;', '{ int DONOTUSE; } NoOp;', ffi_text)
    ffi_text = re.sub('__attribute__\s*\(.*\)', '', ffi_text)
    ffi_text = re.sub('__declspec\s*\([^\)]*\)', '', ffi_text)
    ffi_text = ffi_text.replace('__const', 'const')
    ffi_text = ffi_text.replace('__inline', '')
    ffi_text = ffi_text.replace('__w64', '')
    ffi_text = ffi_text.replace('__cdecl', '')
    ffi_text = ffi_text.replace('__int64', 'long')
    ffi_lines = ffi_text.split(linesep)

    good = find_good_scan([], ffi_lines)
    remaining = ffi_lines[len(good)+1:]

    while len(remaining) > 1:
        l.debug("%d uncertain lines remaining", len(remaining))
        new_good = find_good_scan(good, remaining[1:])
        good += new_good
        remaining = remaining[len(new_good)+1:]

    good += [ 'extern VexControl vex_control;' ]

    open('pyvex/vex_ffi.py', 'w').write('ffi_str = """' + '\n'.join(good) + '"""')

if __name__ == '__main__':
    import sys
    logging.basicConfig(level=logging.DEBUG)
    doit(sys.argv[1])
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright © 2014  Red Hat, Inc.
#
# This copyrighted material is made available to anyone wishing to use,
# modify, copy, or redistribute it subject to the terms and conditions
# of the GNU General Public License v.2, or (at your option) any later
# version.  This program is distributed in the hope that it will be
# useful, but WITHOUT ANY WARRANTY expressed or implied, including the
# implied warranties of MERCHANTABILITY or FITNESS FOR A PARTICULAR
# PURPOSE.  See the GNU General Public License for more details.  You
# should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# Any Red Hat trademarks that are incorporated in the source
# code or documentation are not subject to the GNU General Public
# License and may only be used or replicated with the express permission
# of Red Hat, Inc.
#

import os
import hashlib
import cPickle as pickle

from IPy import IP
import pprint
import dns.resolver

import mirrormanager2.lib

global_caches = dict(
    # key is directoryname
    mirrorlist_cache = {},
    # key is strings in tuple (repo.prefix, arch)
    repo_arch_to_directoryname = {},
    # key is an IPy.IP structure, value is list of host ids
    host_netblock_cache = {},
    # key is hostid, value is list of countries to allow
    host_country_allowed_cache = {},
    host_country_cache = {},
    host_bandwidth_cache = {},
    host_asn_cache = {},
    )

def parent_dir(path):
    return os.path.dirname(path)

def append_value_to_cache(cache, key, value):
    if key not in cache:
        cache[key] = [value]
    else:
        cache[key].append(value)
    return cache

def add_host_to_set(s, hostid):
    s.add(hostid)

def shrink(mc):
    pp = pprint.PrettyPrinter()
    subcaches = ('global', 'byCountry', 'byHostId', 'byCountryInternet2')
    matches = {}
    for d in mc:
        for subcache in subcaches:
            c = mc[d][subcache]
            s = hashlib.sha1(pp.pformat(c)).hexdigest()
            if s in matches:
                mc[d][subcache] = matches[s]
            else:
                matches[s] = c
    return mc

def query_directory_exclusive_host(session):
    table = mirrormanager2.lib.get_directory_exclusive_host(session)
    cache = {}
    for (dname, hostid) in table:
        if dname not in cache:
            cache[dname] = set([hostid])
        else:
            cache[dname].add(hostid)
    return cache

def populate_directory_cache(session):
    global global_caches
    result = mirrormanager2.lib.query_directories(session)

    directory_exclusive_hosts = query_directory_exclusive_host(session)

    def setup_directory_repo_cache(session):
        cache = {}
        for r in mirrormanager2.lib.get_repositories(session):
            if r.directory and r.version and r.arch:
                cache[r.directory.id] = r
        return cache
    def setup_version_ordered_mirrorlist_cache(session):
        cache = {}
        for v in mirrormanager2.lib.get_versions(session):
            cache[v.id] = v.ordered_mirrorlist
        return cache
    def setup_category_topdir_cache(session):
        cache = {}
        for c in mirrormanager2.lib.get_categories(session):
            cache[c.id] = len(c.topdir.name) + 1  # include trailing /
        return cache

    directory_repo_cache = setup_directory_repo_cache(session)
    cat_dirs = mirrormanager2.lib.get_category_directory(session)
    directory_category_cache = {}
    for catdir in cat_dirs:
        append_value_to_cache(
            directory_category_cache, catdir.directory_id, catdir.category_id)
    version_ordered_mirrorlist_cache = setup_version_ordered_mirrorlist_cache(
        session)
    category_topdir_cache = setup_category_topdir_cache(session)

    cache = {}
    for (directory_id, directoryname, hostid, country, hcurl,
            siteprivate, hostprivate, i2, i2_clients) in result:
        if directoryname in directory_exclusive_hosts and \
                hostid not in directory_exclusive_hosts[directoryname]:
            continue

        if directoryname not in cache:
            cache[directoryname] = {
                'global': set(),
                'byCountry': {},
                'byHostId': {},
                'ordered_mirrorlist': True,
                'byCountryInternet2': {}
            }

            repo = directory_repo_cache.get(directory_id)

            if repo is not None and repo.arch is not None:
                global_caches['repo_arch_to_directoryname'][
                    (repo.prefix, repo.arch.name)] = directoryname
                # WARNING - this is a query # fixme use cache
                cache[directoryname][
                    'ordered_mirrorlist'] = repo.version.ordered_mirrorlist

            numcats = len(directory_category_cache.get(directory_id, []))
            if numcats == 0:
                # no category, so we can't know a mirror host's URLs.
                # nothing to add.
                continue
            elif numcats >= 1:
                # any of them will do, so just look at the first one
                category_id = directory_category_cache[directory_id][0]

            cache[directoryname]['subpath'] = directoryname[
                category_topdir_cache[category_id]:]
            del repo

        if country is not None:
            country = country.upper()

        if not siteprivate and not hostprivate:
            add_host_to_set(cache[directoryname]['global'], hostid)

            if country is not None:
                if country not in cache[directoryname]['byCountry']:
                    cache[directoryname]['byCountry'][country] = set()
                add_host_to_set(
                    cache[directoryname]['byCountry'][country], hostid)

        if country is not None and i2 and \
                ((not siteprivate and not hostprivate) or i2_clients):
            if country not in cache[directoryname]['byCountryInternet2']:
                cache[directoryname]['byCountryInternet2'][country] = set()
            add_host_to_set(
                cache[directoryname]['byCountryInternet2'][country], hostid)

        append_value_to_cache(cache[directoryname]['byHostId'], hostid, hcurl)

    global_caches['mirrorlist_cache'] = shrink(cache)

def name_to_ips(name):
    result=[]
    recordtypes=('A', 'AAAA')
    for r in recordtypes:
        try:
            records = dns.resolver.query(name, r)
            for rdata in records:
                try:
                    ip = IP(str(rdata))
                    result.append(ip)
                except:
                    continue
        except:
            continue
    return result

def populate_netblock_cache(cache, host):
    if host.is_active() and len(host.netblocks) > 0:
        for n in list(host.netblocks):
            try:
                ip = IP(n.netblock)
                ips = [ip]
            except ValueError:
                # probably a string
                ips = name_to_ips(n.netblock)

            for ip in ips:
                append_value_to_cache(cache, ip, host.id)
    return cache

def populate_host_country_allowed_cache(cache, host):
    if host.is_active() and len(host.countries_allowed) > 0:
        cache[host.id] = [
            c.country.upper()
            for c in list(host.countries_allowed)
        ]
    return cache

def populate_host_max_connections_cache(cache, host):
    cache[host.id] = host.max_connections
    return cache

def populate_host_bandwidth_cache(cache, host):
    try:
        i = int(host.bandwidth_int)
        if i < 1: i = 1
        elif i > 100000: i = 100000 # max bandwidth 100Gb
        cache[host.id] = i
    except:
        cache[host.id] = 1

    return cache

def populate_host_country_cache(cache, host):
    cache[host.id] = host.country
    return cache

def populate_host_asn_cache(cache, host):
    if not host.asn_clients: return cache
    if host.asn is not None:
        append_value_to_cache(cache, host.asn, host.id)

    for peer_asn in list(host.peer_asns):
        append_value_to_cache(cache, peer_asn.asn, host.id)
    return cache

def repository_redirect_cache(session):
    cache = {}
    for r in mirrormanager2.lib.get_reporedirect(session):
        cache[r.from_repo] = r.to_repo
    return cache

def country_continent_redirect_cache(session):
    cache = {}
    for c in mirrormanager2.lib.get_country_continent_redirect(session):
        cache[c.country] = c.continent
    return cache

def disabled_repository_cache(session):
    cache = {}
    for r in mirrormanager2.lib.get_repositories(session):
        if r.disabled:
            cache[r.prefix] = True
    return cache

def file_details_cache(session):
    # cache{directoryname}{filename}[{details}]
    cache = {}
    # materialize this select to avoid making hundreds of thousands of queries
    for d in mirrormanager2.lib.get_directories(session):
        if len(d.fileDetails) > 0:
            cache[d.name] = {}
            for fd in list(d.fileDetails):
                details = dict(
                    timestamp=fd.timestamp,
                    sha1=fd.sha1,
                    md5=fd.md5,
                    sha256=fd.sha256,
                    sha512=fd.sha512,
                    size=fd.size)
                append_value_to_cache(cache[d.name], fd.filename, details)
    return cache

def hcurl_cache(session):
    cache = {}
    for hcurl in mirrormanager2.lib.get_host_category_url(session):
        cache[hcurl.id] = hcurl.url
    return cache

def location_cache(session):
    cache = {}
    for location in mirrormanager2.lib.get_locations(session):
        cache[location.name] = [host.id for host in list(location.hosts)]
    return cache

def netblock_country_cache(session):
    cache = {}
    for n in mirrormanager2.lib.get_netblock_country(session):
        ip = None
        try:
            ip = IP(n.netblock)
        except:
            continue
        # guaranteed to be unique by database constraints
        cache[ip] = (n.country)

    return cache

def populate_host_caches(session):
    n = dict()
    ca = dict()
    b = dict()
    cc = dict()
    a = dict()
    mc = dict()

    for host in list(mirrormanager2.lib.get_hosts(session)):
        n = populate_netblock_cache(n, host)
        ca = populate_host_country_allowed_cache(ca, host)
        b = populate_host_bandwidth_cache(b, host)
        cc = populate_host_country_cache(cc, host)
        a = populate_host_asn_cache(a, host)
        mc = populate_host_max_connections_cache(mc, host)

    global global_caches
    global_caches['host_netblock_cache'] = n
    global_caches['host_country_allowed_cache'] = ca
    global_caches['host_bandwidth_cache'] = b
    global_caches['host_country_cache'] = cc
    global_caches['host_asn_cache'] = a
    global_caches['host_max_connections_cache'] = mc

def populate_all_caches(session):
    populate_host_caches(session)
    populate_directory_cache(session)

def dump_caches(session, filename):
    data = {
        'mirrorlist_cache': global_caches['mirrorlist_cache'],
        'host_netblock_cache': global_caches['host_netblock_cache'],
        'host_country_allowed_cache': global_caches['host_country_allowed_cache'],
        'host_bandwidth_cache': global_caches['host_bandwidth_cache'],
        'host_country_cache': global_caches['host_country_cache'],
        'host_max_connections_cache': global_caches['host_max_connections_cache'],
        'asn_host_cache': global_caches['host_asn_cache'], # yeah I misnamed this
        'repo_arch_to_directoryname': global_caches['repo_arch_to_directoryname'],
        'repo_redirect_cache': repository_redirect_cache(session),
        'country_continent_redirect_cache': country_continent_redirect_cache(session),
        'disabled_repositories': disabled_repository_cache(session),
        'file_details_cache': file_details_cache(session),
        'hcurl_cache': hcurl_cache(session),
        'location_cache': location_cache(session),
        'netblock_country_cache': netblock_country_cache(session),
    }

    try:
        f = open(filename, 'w')
        pickle.dump(data, f)
        f.close()
        print 'Pickle generated at %s' % filename
    except Exception as err:
        print 'Error generating the pickle (%s): %s' % (
            filename, err)
        pass
<commit_msg>Style fix<commit_after># -*- coding: utf-8 -*-
#
# Copyright © 2014  Red Hat, Inc.
#
# This copyrighted material is made available to anyone wishing to use,
# modify, copy, or redistribute it subject to the terms and conditions
# of the GNU General Public License v.2, or (at your option) any later
# version.  This program is distributed in the hope that it will be
# useful, but WITHOUT ANY WARRANTY expressed or implied, including the
# implied warranties of MERCHANTABILITY or FITNESS FOR A PARTICULAR
# PURPOSE.  See the GNU General Public License for more details.  You
# should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# Any Red Hat trademarks that are incorporated in the source
# code or documentation are not subject to the GNU General Public
# License and may only be used or replicated with the express permission
# of Red Hat, Inc.
#

import os
import hashlib
import cPickle as pickle

from IPy import IP
import pprint
import dns.resolver

import mirrormanager2.lib

global_caches = dict(
    # key is directoryname
    mirrorlist_cache = {},
    # key is strings in tuple (repo.prefix, arch)
    repo_arch_to_directoryname = {},
    # key is an IPy.IP structure, value is list of host ids
    host_netblock_cache = {},
    # key is hostid, value is list of countries to allow
    host_country_allowed_cache = {},
    host_country_cache = {},
    host_bandwidth_cache = {},
    host_asn_cache = {},
    )

def parent_dir(path):
    return os.path.dirname(path)

def append_value_to_cache(cache, key, value):
    if key not in cache:
        cache[key] = [value]
    else:
        cache[key].append(value)
    return cache

def add_host_to_set(s, hostid):
    s.add(hostid)

def shrink(mc):
    pp = pprint.PrettyPrinter()
    subcaches = ('global', 'byCountry', 'byHostId', 'byCountryInternet2')
    matches = {}
    for d in mc:
        for subcache in subcaches:
            c = mc[d][subcache]
            s = hashlib.sha1(pp.pformat(c)).hexdigest()
            if s in matches:
                mc[d][subcache] = matches[s]
            else:
                matches[s] = c
    return mc

def query_directory_exclusive_host(session):
    table = mirrormanager2.lib.get_directory_exclusive_host(session)
    cache = {}
    for (dname, hostid) in table:
        if dname not in cache:
            cache[dname] = set([hostid])
        else:
            cache[dname].add(hostid)
    return cache

def populate_directory_cache(session):
    global global_caches
    result = mirrormanager2.lib.query_directories(session)

    directory_exclusive_hosts = query_directory_exclusive_host(session)

    def setup_directory_repo_cache(session):
        cache = {}
        for r in mirrormanager2.lib.get_repositories(session):
            if r.directory and r.version and r.arch:
                cache[r.directory.id] = r
        return cache
    def setup_version_ordered_mirrorlist_cache(session):
        cache = {}
        for v in mirrormanager2.lib.get_versions(session):
            cache[v.id] = v.ordered_mirrorlist
        return cache
    def setup_category_topdir_cache(session):
        cache = {}
        for c in mirrormanager2.lib.get_categories(session):
            cache[c.id] = len(c.topdir.name) + 1  # include trailing /
        return cache

    directory_repo_cache = setup_directory_repo_cache(session)
    cat_dirs = mirrormanager2.lib.get_category_directory(session)
    directory_category_cache = {}
    for catdir in cat_dirs:
        append_value_to_cache(
            directory_category_cache, catdir.directory_id, catdir.category_id)
    version_ordered_mirrorlist_cache = setup_version_ordered_mirrorlist_cache(
        session)
    category_topdir_cache = setup_category_topdir_cache(session)

    cache = {}
    for (directory_id, directoryname, hostid, country, hcurl,
            siteprivate, hostprivate, i2, i2_clients) in result:
        if directoryname in directory_exclusive_hosts and \
                hostid not in directory_exclusive_hosts[directoryname]:
            continue

        if directoryname not in cache:
            cache[directoryname] = {
                'global': set(),
                'byCountry': {},
                'byHostId': {},
                'ordered_mirrorlist': True,
                'byCountryInternet2': {}
            }

            repo = directory_repo_cache.get(directory_id)

            if repo is not None and repo.arch is not None:
                global_caches['repo_arch_to_directoryname'][
                    (repo.prefix, repo.arch.name)] = directoryname
                # WARNING - this is a query # fixme use cache
                cache[directoryname][
                    'ordered_mirrorlist'] = repo.version.ordered_mirrorlist

            numcats = len(directory_category_cache.get(directory_id, []))
            if numcats == 0:
                # no category, so we can't know a mirror host's URLs.
                # nothing to add.
                continue
            elif numcats >= 1:
                # any of them will do, so just look at the first one
                category_id = directory_category_cache[directory_id][0]

            cache[directoryname]['subpath'] = directoryname[
                category_topdir_cache[category_id]:]
            del repo

        if country is not None:
            country = country.upper()

        if not siteprivate and not hostprivate:
            add_host_to_set(cache[directoryname]['global'], hostid)

            if country is not None:
                if country not in cache[directoryname]['byCountry']:
                    cache[directoryname]['byCountry'][country] = set()
                add_host_to_set(
                    cache[directoryname]['byCountry'][country], hostid)

        if country is not None and i2 and \
                ((not siteprivate and not hostprivate) or i2_clients):
            if country not in cache[directoryname]['byCountryInternet2']:
                cache[directoryname]['byCountryInternet2'][country] = set()
            add_host_to_set(
                cache[directoryname]['byCountryInternet2'][country], hostid)

        append_value_to_cache(cache[directoryname]['byHostId'], hostid, hcurl)

    global_caches['mirrorlist_cache'] = shrink(cache)

def name_to_ips(name):
    result=[]
    recordtypes=('A', 'AAAA')
    for r in recordtypes:
        try:
            records = dns.resolver.query(name, r)
            for rdata in records:
                try:
                    ip = IP(str(rdata))
                    result.append(ip)
                except:
                    continue
        except:
            continue
    return result

def populate_netblock_cache(cache, host):
    if host.is_active() and len(host.netblocks) > 0:
        for n in list(host.netblocks):
            try:
                ip = IP(n.netblock)
                ips = [ip]
            except ValueError:
                # probably a string
                ips = name_to_ips(n.netblock)

            for ip in ips:
                append_value_to_cache(cache, ip, host.id)
    return cache

def populate_host_country_allowed_cache(cache, host):
    if host.is_active() and len(host.countries_allowed) > 0:
        cache[host.id] = [
            c.country.upper()
            for c in list(host.countries_allowed)
        ]
    return cache

def populate_host_max_connections_cache(cache, host):
    cache[host.id] = host.max_connections
    return cache

def populate_host_bandwidth_cache(cache, host):
    try:
        i = int(host.bandwidth_int)
        if i < 1: i = 1
        elif i > 100000: i = 100000 # max bandwidth 100Gb
        cache[host.id] = i
    except:
        cache[host.id] = 1

    return cache

def populate_host_country_cache(cache, host):
    cache[host.id] = host.country
    return cache

def populate_host_asn_cache(cache, host):
    if not host.asn_clients:
        return cache

    if host.asn is not None:
        append_value_to_cache(cache, host.asn, host.id)

    for peer_asn in list(host.peer_asns):
        append_value_to_cache(cache, peer_asn.asn, host.id)
    return cache

def repository_redirect_cache(session):
    cache = {}
    for r in mirrormanager2.lib.get_reporedirect(session):
        cache[r.from_repo] = r.to_repo
    return cache

def country_continent_redirect_cache(session):
    cache = {}
    for c in mirrormanager2.lib.get_country_continent_redirect(session):
        cache[c.country] = c.continent
    return cache

def disabled_repository_cache(session):
    cache = {}
    for r in mirrormanager2.lib.get_repositories(session):
        if r.disabled:
            cache[r.prefix] = True
    return cache

def file_details_cache(session):
    # cache{directoryname}{filename}[{details}]
    cache = {}
    # materialize this select to avoid making hundreds of thousands of queries
    for d in mirrormanager2.lib.get_directories(session):
        if len(d.fileDetails) > 0:
            cache[d.name] = {}
            for fd in list(d.fileDetails):
                details = dict(
                    timestamp=fd.timestamp,
                    sha1=fd.sha1,
                    md5=fd.md5,
                    sha256=fd.sha256,
                    sha512=fd.sha512,
                    size=fd.size)
                append_value_to_cache(cache[d.name], fd.filename, details)
    return cache

def hcurl_cache(session):
    cache = {}
    for hcurl in mirrormanager2.lib.get_host_category_url(session):
        cache[hcurl.id] = hcurl.url
    return cache

def location_cache(session):
    cache = {}
    for location in mirrormanager2.lib.get_locations(session):
        cache[location.name] = [host.id for host in list(location.hosts)]
    return cache

def netblock_country_cache(session):
    cache = {}
    for n in mirrormanager2.lib.get_netblock_country(session):
        ip = None
        try:
            ip = IP(n.netblock)
        except:
            continue
        # guaranteed to be unique by database constraints
        cache[ip] = (n.country)

    return cache

def populate_host_caches(session):
    n = dict()
    ca = dict()
    b = dict()
    cc = dict()
    a = dict()
    mc = dict()

    for host in list(mirrormanager2.lib.get_hosts(session)):
        n = populate_netblock_cache(n, host)
        ca = populate_host_country_allowed_cache(ca, host)
        b = populate_host_bandwidth_cache(b, host)
        cc = populate_host_country_cache(cc, host)
        a = populate_host_asn_cache(a, host)
        mc = populate_host_max_connections_cache(mc, host)

    global global_caches
    global_caches['host_netblock_cache'] = n
    global_caches['host_country_allowed_cache'] = ca
    global_caches['host_bandwidth_cache'] = b
    global_caches['host_country_cache'] = cc
    global_caches['host_asn_cache'] = a
    global_caches['host_max_connections_cache'] = mc

def populate_all_caches(session):
    populate_host_caches(session)
    populate_directory_cache(session)

def dump_caches(session, filename):
    data = {
        'mirrorlist_cache': global_caches['mirrorlist_cache'],
        'host_netblock_cache': global_caches['host_netblock_cache'],
        'host_country_allowed_cache': global_caches['host_country_allowed_cache'],
        'host_bandwidth_cache': global_caches['host_bandwidth_cache'],
        'host_country_cache': global_caches['host_country_cache'],
        'host_max_connections_cache': global_caches['host_max_connections_cache'],
        'asn_host_cache': global_caches['host_asn_cache'], # yeah I misnamed this
        'repo_arch_to_directoryname': global_caches['repo_arch_to_directoryname'],
        'repo_redirect_cache': repository_redirect_cache(session),
        'country_continent_redirect_cache': country_continent_redirect_cache(session),
        'disabled_repositories': disabled_repository_cache(session),
        'file_details_cache': file_details_cache(session),
        'hcurl_cache': hcurl_cache(session),
        'location_cache': location_cache(session),
        'netblock_country_cache': netblock_country_cache(session),
    }

    try:
        f = open(filename, 'w')
        pickle.dump(data, f)
        f.close()
        print 'Pickle generated at %s' % filename
    except Exception as err:
        print 'Error generating the pickle (%s): %s' % (
            filename, err)
        pass
<|endoftext|>


<commit_before>from flask import Flask, redirect, url_for, session, request
from flask_oauth import OAuth
from config import GOOGLE_OAUTH_ID, GOOGLE_OAUTH_SECRET, \
                    TWITTER_OAUTH_ID, TWITTER_OAUTH_SECRET, \
                    FACEBOOK_OAUTH_ID, FACEBOOK_OAUTH_SECRET

oauth = OAuth()

google = oauth.remote_app('google',
    base_url='https://www.google.com/accounts/',
    authorize_url='https://accounts.google.com/o/oauth2/auth',
    request_token_url=None,
    request_token_params={'scope': 'https://www.googleapis.com/auth/userinfo.email https://www.googleapis.com/auth/userinfo.profile',
        'response_type': 'code'},
    access_token_url='https://accounts.google.com/o/oauth2/token',
    access_token_method='POST',
    access_token_params={'grant_type': 'authorization_code'},
    consumer_key=GOOGLE_OAUTH_ID,
    consumer_secret=GOOGLE_OAUTH_SECRET)

# Use Twitter as example remote application
twitter = oauth.remote_app('twitter',
    # unless absolute urls are used to make requests, this will be added
    # before all URLs.  This is also true for request_token_url and others.
    base_url='https://api.twitter.com/1.1/',
    # where flask should look for new request tokens
    request_token_url='https://api.twitter.com/oauth/request_token',
    # where flask should exchange the token with the remote application
    access_token_url='https://api.twitter.com/oauth/access_token',
    # twitter knows two authorizatiom URLs.  /authorize and /authenticate.
    # they mostly work the same, but for sign on /authenticate is
    # expected because this will give the user a slightly different
    # user interface on the twitter side.
    authorize_url='https://api.twitter.com/oauth/authenticate',
    # the consumer keys from the twitter application registry.
    consumer_key=TWITTER_OAUTH_ID,
    consumer_secret=TWITTER_OAUTH_SECRET)

facebook = oauth.remote_app('facebook',
    base_url='https://graph.facebook.com/',
    request_token_url=None,
    access_token_url='/oauth/access_token',
    authorize_url='https://www.facebook.com/dialog/oauth',
    consumer_key=FACEBOOK_OAUTH_ID,
    consumer_secret=FACEBOOK_OAUTH_SECRET,
    request_token_params={'scope': 'email'})

<commit_msg>Removes unused imports.<commit_after>from flask_oauth import OAuth
from config import GOOGLE_OAUTH_ID, GOOGLE_OAUTH_SECRET, TWITTER_OAUTH_ID, TWITTER_OAUTH_SECRET, FACEBOOK_OAUTH_ID, FACEBOOK_OAUTH_SECRET

oauth = OAuth()

google = oauth.remote_app('google',
    base_url='https://www.google.com/accounts/',
    authorize_url='https://accounts.google.com/o/oauth2/auth',
    request_token_url=None,
    request_token_params={'scope': 'https://www.googleapis.com/auth/userinfo.email https://www.googleapis.com/auth/userinfo.profile',
        'response_type': 'code'},
    access_token_url='https://accounts.google.com/o/oauth2/token',
    access_token_method='POST',
    access_token_params={'grant_type': 'authorization_code'},
    consumer_key=GOOGLE_OAUTH_ID,
    consumer_secret=GOOGLE_OAUTH_SECRET)

# Use Twitter as example remote application
twitter = oauth.remote_app('twitter',
    # unless absolute urls are used to make requests, this will be added
    # before all URLs.  This is also true for request_token_url and others.
    base_url='https://api.twitter.com/1.1/',
    # where flask should look for new request tokens
    request_token_url='https://api.twitter.com/oauth/request_token',
    # where flask should exchange the token with the remote application
    access_token_url='https://api.twitter.com/oauth/access_token',
    # twitter knows two authorizatiom URLs.  /authorize and /authenticate.
    # they mostly work the same, but for sign on /authenticate is
    # expected because this will give the user a slightly different
    # user interface on the twitter side.
    authorize_url='https://api.twitter.com/oauth/authenticate',
    # the consumer keys from the twitter application registry.
    consumer_key=TWITTER_OAUTH_ID,
    consumer_secret=TWITTER_OAUTH_SECRET)

facebook = oauth.remote_app('facebook',
    base_url='https://graph.facebook.com/',
    request_token_url=None,
    access_token_url='/oauth/access_token',
    authorize_url='https://www.facebook.com/dialog/oauth',
    consumer_key=FACEBOOK_OAUTH_ID,
    consumer_secret=FACEBOOK_OAUTH_SECRET,
    request_token_params={'scope': 'email'})
<|endoftext|>


<commit_before># (C) Datadog, Inc. 2019-present
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)
import pytest

from .metrics import ALL_METRICS

pytestmark = pytest.mark.e2e

def test_check(dd_agent_check, instance):
    # We do not do aggregator.assert_all_metrics_covered() because depending on timing, some other metrics may appear
    aggregator = dd_agent_check(instance, rate=True)

    server_tag = 'server:{}'.format(instance['server'])
    port_tag = 'port:{}'.format(instance['port'])
    for metric in ALL_METRICS:
        aggregator.assert_metric_has_tag(metric, server_tag)
        aggregator.assert_metric_has_tag(metric, port_tag)
        aggregator.assert_metric_has_tag(metric, 'db:default')
        aggregator.assert_metric_has_tag(metric, 'foo:bar')

    aggregator.assert_metric('clickhouse.table.replicated.total', 2)
    aggregator.assert_metric(
        'clickhouse.dictionary.item.current', tags=[server_tag, port_tag, 'db:default', 'foo:bar', 'dictionary:test']
    )
<commit_msg>Update test_e2e.py (#6355)<commit_after># (C) Datadog, Inc. 2019-present
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)
import pytest

from .metrics import ALL_METRICS

pytestmark = pytest.mark.e2e

def test_check(dd_agent_check, instance):
    # We do not do aggregator.assert_all_metrics_covered() because depending on timing, some other metrics may appear
    aggregator = dd_agent_check(instance, rate=True)

    server_tag = 'server:{}'.format(instance['server'])
    port_tag = 'port:{}'.format(instance['port'])
    for metric in ALL_METRICS:
        aggregator.assert_metric_has_tag(metric, server_tag)
        aggregator.assert_metric_has_tag(metric, port_tag)
        aggregator.assert_metric_has_tag(metric, 'db:default')
        aggregator.assert_metric_has_tag(metric, 'foo:bar')

    aggregator.assert_metric('clickhouse.table.replicated.total')
    aggregator.assert_metric(
        'clickhouse.dictionary.item.current', tags=[server_tag, port_tag, 'db:default', 'foo:bar', 'dictionary:test']
    )
<|endoftext|>