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

__author__ = "Danny Broberg, Shyam Dwaraknath, Bharat Medasani, Nils Zimmermann, Geoffroy Hautier"
__copyright__ = "Copyright 2014, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Danny Broberg, Shyam Dwaraknath"
__email__ = "dbroberg@berkeley.edu, shyamd@lbl.gov"
__status__ = "Development"
__date__ = "January 11, 2018"

from math import exp
import numpy as np
from monty.json import MSONable
from scipy import integrate
from scipy.spatial import HalfspaceIntersection
from scipy.optimize import bisect
from itertools import groupby, chain

from pymatgen.analysis.defects.utils import kb

class DefectPhaseDiagram(MSONable):
    """
    This is similar to a PhaseDiagram object in pymatgen, but has ability to do quick analysis of defect formation energies
    when fed DefectEntry objects

    uses many of the capabilities from PyCDT's DefectsAnalyzer class...

    This class is able to get:
        a) stability of charge states for a given defect,
        b) list of all formation ens
        c) transition levels in the gap
        d)

    Args:
        dentries ([DefectEntry]): A list of DefectEntry objects
    """

    def __init__(self, entries, vbm, band_gap, filter_compatible=True):
        self.vbm = vbm
        self.band_gap = band_gap
        self.filter_compatible = filter_compatible

        if filter_compatible:
            self.entries = [e for e in entries if e.parameters.get("is_compatible", True)]
        else:
            self.entries = entries

        self.find_stable_charges()

    def find_stable_charges(self):
        """
        Sets the stable charges and transition states for a series of
        defect entries. This function uses scipy's HalfspaceInterection
        to oncstruct the polygons corresponding to defect stability as
        a function of the Fermi-level. The Halfspace Intersection
        constructs N-dimensional hyperplanes, in this case N=2,  based
        on the equation of defect formation energy with considering chemical
        potentials:
            E_form = E_0^{Corrected} + Q_{defect}*(E_{VBM} + E_{Fermi})

        Extra hyperplanes are constructed to bound this space so that
        the algorithm can actually find enclosed region.

        This code was modeled after the Halfspace Intersection code for
        the Pourbaix Diagram
        """

        def similar_defect(a):
            """
            Used to filter out similar defects of different charges which
            are defined by the same type and location
            """
            return (a.name, a.site)

        # Limits for search
        # E_fermi = { -1 eV to band gap+1}
        # the 1 eV padding provides
        # E_formation. = { -21 eV to 20 eV}
        limits = [[-1, self.band_gap + 1], [-21, 20]]

        stable_entries = {}
        stable_charges = {}
        finished_charges = {}
        transition_level_map = {}

        # Grouping by defect types
        for _, defects in groupby(sorted(self.entries, key=similar_defect), similar_defect):
            defects = list(defects)

            # prepping coefficient matrix forx half-space intersection
            # [-Q, 1, -1*(E_form+Q*VBM)] -> -Q*E_fermi+E+-1*(E_form+Q*VBM) <= 0  where E_fermi and E are the variables in the hyperplanes
            hyperplanes = np.array(
                [[-1.0 * entry.charge, 1, -1.0 * (entry.energy + entry.charge * self.vbm)] for entry in defects])

            border_hyperplanes = [[-1, 0, limits[0][0]], [1, 0, -1 * limits[0][1]], [0, -1, limits[1][0]],
                                  [0, 1, -1 * limits[1][1]]]
            hs_hyperplanes = np.vstack([hyperplanes, border_hyperplanes])

            interior_point = [self.band_gap / 2, -20]

            hs_ints = HalfspaceIntersection(hs_hyperplanes, np.array(interior_point))

            # Group the intersections and coresponding facets
            ints_and_facets = zip(hs_ints.intersections, hs_ints.dual_facets)
            # Only inlcude the facets corresponding to entries, not the boundaries
            total_entries = len(defects)
            ints_and_facets = filter(lambda int_and_facet: all(np.array(int_and_facet[1]) < total_entries),
                                     ints_and_facets)
            # sort based on transition level
            ints_and_facets = list(sorted(ints_and_facets, key=lambda int_and_facet: int_and_facet[0][0]))
            # Unpack into lists
            intersections, facets = zip(*ints_and_facets)
            # Map of transition level: charge states

            transition_level_map[defects[0].name] = {
                intersection[0]: [defects[i].charge for i in facet]
                for intersection, facet in ints_and_facets
            }

            stable_entries[defects[0].name] = [defects[i] for dual in facets for i in dual]

            finished_charges[defects[0].name] = [defect.charge for defect in defects]

        self.transition_level_map = transition_level_map
        self.transition_levels = {
            defect_name: list(defect_tls.keys())
            for defect_name, defect_tls in transition_level_map.items()
        }
        self.stable_entries = stable_entries
        self.finished_charges = finished_charges
        self.stable_charges = {
            defect_name: [entry.charge for entry in entries]
            for defect_name, entries in stable_entries.items()
        }

    def OLDfind_stable_charges(self):
        """
        Set the DefectPhaseDiagram attributes (which dont depend on chemical potential)

        THIS is a wrapper since find_stable_charges does not currently exist...will switch over when able
        """
        def similar_defect(a, b):
            """
            Used to filter out similar defects of different charges which
            are defined by the same type and location
            """
            return a.name == b.name and a.site == b.site

        self.stable_charges = {} #keys are defect names, items are list of charge states that are stable
        self.finished_charges = {} #keys are defect names, items are list of charge states that are included in the phase diagram
        self.transition_levels = {} #keys are defect names, items are list of [fermi level for transition, previous q, next q] sets
        xlim = (-0.1, self.band_gap+.1)
        nb_steps = 20 #SUPER bare bones because it is so slow...
        x = np.arange(xlim[0], xlim[1], (xlim[1]-xlim[0])/nb_steps)
        list_elts = [elt  for dfct in self.entries  for elt in dfct.generate_defect_structure().composition.elements]
        no_chempots = {elt: 0. for elt in set(list_elts)} #zerod chemical potentials for calculating stable defects

        # Grouping by defect types
        for _, defects in groupby(sorted(self.entries, key=similar_defect), similar_defect):
            defects = list(defects)
            t = defects[0].name
            print('Analyzing {}'.format(t))
            trans_level = []
            chg_type = []
            prev_min_q, cur_min_q = None, None
            for x_step in x:
                miny = 10000
                for dfct in defects:
                    val = dfct.formation_energy(chemical_potentials = no_chempots, fermi_level=x_step)
                    if val < miny:
                        miny = val
                        cur_min_q = dfct.charge

                if prev_min_q is not None:
                    if cur_min_q != prev_min_q:
                        trans_level.append((x_step, prev_min_q, cur_min_q))
                    if cur_min_q not in chg_type:
                        chg_type.append(cur_min_q)
                prev_min_q = cur_min_q

            self.stable_charges[t] = chg_type[:]
            self.finished_charges[t] = [e.charge for e in self.entries if e.name == t]
            self.transition_levels[t] = trans_level[:]

    @property
    def defect_types(self):
        """
        List types of defects existing in the DefectPhaseDiagram
        """
        return list(self.finished_charges.keys())

    @property
    def all_stable_entries(self):
        """
        List all stable entries (defect+charge) in the DefectPhaseDiagram
        """
        return set(chain.from_iterable(self.stable_entries.values()))

    @property
    def all_unstable_entries(self):
        """
        List all unstable entries (defect+charge) in the DefectPhaseDiagram
        """
        all_stable_entries = self.all_stable_entries
        return [e for e in self.entries if e not in all_stable_entries]

    def defect_concentrations(self, chemical_potentials, temperature=300, fermi_level=0.):
        """
        Give list of all concentrations at specified efermi in the DefectPhaseDiagram
        args:
            chemical_potentials = {Element: number} is dictionary of chemical potentials to provide formation energies for
            temperature = temperature to produce concentrations from
            fermi_level: (float) is fermi level relative to valence band maximum
                Default efermi = 0 = VBM energy
        returns:
            list of dictionaries of defect concentrations
        """
        concentrations = []
        for dfct in self.all_stable_entries:
            concentrations.append({
                'conc':
                dfct.defect_concentration(
                    chemical_potentials=chemical_potentials, temperature=temperature, fermi_level=fermi_level),
                'name':
                dfct.name,
                'charge':
                dfct.charge
            })

            # TODO: Should be this entry: concentration?
        return concentrations

    def suggest_charges(self, tolerance=0.1):
        """
        Suggest possible charges for defects to computee based on proximity
        of known transitions from entires to VBM and CBM

        Args:
            tolerance (float): tolerance with respect to the VBM and CBM to
                    `          continue to compute new charges
        """
        recommendations = {}

        for def_type in self.defect_types:
            test_charges = np.arange(
                np.min(self.stable_charges[def_type]) - 1,
                np.max(self.stable_charges[def_type]) + 2)
            test_charges = [charge for charge in test_charges if charge not in self.finished_charges[def_type]]

            # More positive charges will shift the minimum transition level down
            # Max charge is limited by this if its transition level is close to VBM
            min_tl = min(self.transition_level_map[def_type].keys())
            if min_tl < tolerance:
                max_charge = max(self.transition_level_map[def_type][min_tl])
                test_charges = [charge for charge in test_charges if charge < max_charge]

            # More negative charges will shift the maximum transition level up
            # Minimum charge is limited by this if transition level is near CBM
            max_tl = max(self.transition_level_map[def_type].keys())
            if max_tl > (self.band_gap - tolerance):
                min_charge = min(self.transition_level_map[def_type][max_tl])
                test_charges = [charge for charge in test_charges if charge > min_charge]

            recommendations[def_type] = test_charges

        return recommendations

    def solve_for_fermi_energy(self, temperature, chemical_potentials, bulk_dos):
        """
        Solve for the Fermi energy self-consistently as a function of T
        and p_O2
        Observations are Defect concentrations, electron and hole conc
        Args:
            bulk_dos: bulk system dos (pymatgen Dos object)
            gap: Can be used to specify experimental gap.
                Will be useful if the self consistent Fermi level
                is > DFT gap
        Returns:
            Fermi energy
        """

        fdos = FermiDos(bulk_dos, bandgap=self.band_gap)

        def _get_total_q(ef):

            qd_tot = sum([
                d['charge'] * d['conc']
                for d in self.list_defect_concentrations(
                    chemical_potentials=chemical_potentials, temperature=temperature, fermi_level=ef)
            ])
            qd_tot += fdos.get_doping(fermi=ef, T=temperature)
            return qd_tot

        return bisect(_get_total_q, -1., self.band_gap + 1.)

    def solve_non_eq_fermilevel(self, lowT, highT, bulk_dos, show_carriers=True, hold_htemp_ef=True):
        """
        Solve for the Fermi energy at a low temperature value, but with defect concentrations frozen in
            at some higher temperature value This means that only the fermi level and free charge carriers concentrations
            will be solved self consistently (simulating a 'frozen-in' defect concentration approach)

        Args:
            lowT: the temperature at which you want to solve for fermi level and free carrier concentrations
            highT: the fixed temperature that sets the defect concentrations at a higher temperature
            bulk_dos: bulk system dos (pymatgen Dos object)
            show_carriers: Dictates the outputs of this function
                Set to True if you want the new carrier concentrations in addition to the fermi levels
            hold_htemp_ef: Specifies whether the fermi level will be fixed for the defect concentrations at high T
                Set to False if you want the fermi level of the high temperature defects to be the same as the fermi level
                for the low temperature charge carriers (see below note on implementation for explanation of this flag)
        Returns:
            Fermi energy

        A NOTE ON IMPLEMENTATION: while testing this function, it became apparent that negative defect formation energies
            can sometimes cause stochastic results for the fermi level as a result of the defect concentrations not
            being allowed to change within the self consistent result. To circumvent this, this function has a flag ('hold_htemp_ef')
            to allow for the fermi level of high temperature defects to vary within the self consistent approach.
            While this is no longer the fully physical 'frozen-in' approach, it helps remove the stocastic results
            that result from numerical issues

        """
        from scipy.optimize import bisect
        q_h_cont = IntrinsicCarrier(bulk_dos, exp_gap=self.band_gap)

        if hold_htemp_ef:
            self.solve_for_fermi_energy(bulk_dos, gap=self.band_gap)
            defectlist = self.get_defects_concentration(highT, self._ef, unitcell=False)

        def noneq_total_q(ef):
            qd_tot = 0
            if hold_htemp_ef:
                for d in defectlist:
                    qd_tot += d['charge'] * d['conc']
            else:
                for d in self.get_defects_concentration(highT, ef, unitcell=False):
                    qd_tot += d['charge'] * d['conc']
            ef_ref_cbm = ef - self.band_gap
            nden = q_h_cont.get_n_density(ef_ref_cbm, lowT, ref='CBM', unitcell=False)
            pden = q_h_cont.get_p_density(ef, lowT, unitcell=False)
            qd_tot += pden - nden
            return qd_tot

        non_eq_ef = bisect(lambda e: noneq_total_q(e), -1., self.band_gap + 1.)
        if show_carriers:
            ef_ref_cbm = non_eq_ef - self.band_gap
            nden = q_h_cont.get_n_density(ef_ref_cbm, lowT, ref='CBM', unitcell=False)
            pden = q_h_cont.get_p_density(non_eq_ef, lowT, unitcell=False)
            return non_eq_ef, nden, pden
        else:
            return non_eq_ef

<commit_msg>another bug in old stable chgs<commit_after>#!/usr/bin/env python

__author__ = "Danny Broberg, Shyam Dwaraknath, Bharat Medasani, Nils Zimmermann, Geoffroy Hautier"
__copyright__ = "Copyright 2014, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Danny Broberg, Shyam Dwaraknath"
__email__ = "dbroberg@berkeley.edu, shyamd@lbl.gov"
__status__ = "Development"
__date__ = "January 11, 2018"

from math import exp
import numpy as np
from monty.json import MSONable
from scipy import integrate
from scipy.spatial import HalfspaceIntersection
from scipy.optimize import bisect
from itertools import groupby, chain

from pymatgen.analysis.defects.utils import kb

class DefectPhaseDiagram(MSONable):
    """
    This is similar to a PhaseDiagram object in pymatgen, but has ability to do quick analysis of defect formation energies
    when fed DefectEntry objects

    uses many of the capabilities from PyCDT's DefectsAnalyzer class...

    This class is able to get:
        a) stability of charge states for a given defect,
        b) list of all formation ens
        c) transition levels in the gap
        d)

    Args:
        dentries ([DefectEntry]): A list of DefectEntry objects
    """

    def __init__(self, entries, vbm, band_gap, filter_compatible=True):
        self.vbm = vbm
        self.band_gap = band_gap
        self.filter_compatible = filter_compatible

        if filter_compatible:
            self.entries = [e for e in entries if e.parameters.get("is_compatible", True)]
        else:
            self.entries = entries

        self.find_stable_charges()

    def find_stable_charges(self):
        """
        Sets the stable charges and transition states for a series of
        defect entries. This function uses scipy's HalfspaceInterection
        to oncstruct the polygons corresponding to defect stability as
        a function of the Fermi-level. The Halfspace Intersection
        constructs N-dimensional hyperplanes, in this case N=2,  based
        on the equation of defect formation energy with considering chemical
        potentials:
            E_form = E_0^{Corrected} + Q_{defect}*(E_{VBM} + E_{Fermi})

        Extra hyperplanes are constructed to bound this space so that
        the algorithm can actually find enclosed region.

        This code was modeled after the Halfspace Intersection code for
        the Pourbaix Diagram
        """

        def similar_defect(a):
            """
            Used to filter out similar defects of different charges which
            are defined by the same type and location
            """
            return (a.name, a.site)

        # Limits for search
        # E_fermi = { -1 eV to band gap+1}
        # the 1 eV padding provides
        # E_formation. = { -21 eV to 20 eV}
        limits = [[-1, self.band_gap + 1], [-21, 20]]

        stable_entries = {}
        stable_charges = {}
        finished_charges = {}
        transition_level_map = {}

        # Grouping by defect types
        for _, defects in groupby(sorted(self.entries, key=similar_defect), similar_defect):
            defects = list(defects)

            # prepping coefficient matrix forx half-space intersection
            # [-Q, 1, -1*(E_form+Q*VBM)] -> -Q*E_fermi+E+-1*(E_form+Q*VBM) <= 0  where E_fermi and E are the variables in the hyperplanes
            hyperplanes = np.array(
                [[-1.0 * entry.charge, 1, -1.0 * (entry.energy + entry.charge * self.vbm)] for entry in defects])

            border_hyperplanes = [[-1, 0, limits[0][0]], [1, 0, -1 * limits[0][1]], [0, -1, limits[1][0]],
                                  [0, 1, -1 * limits[1][1]]]
            hs_hyperplanes = np.vstack([hyperplanes, border_hyperplanes])

            interior_point = [self.band_gap / 2, -20]

            hs_ints = HalfspaceIntersection(hs_hyperplanes, np.array(interior_point))

            # Group the intersections and coresponding facets
            ints_and_facets = zip(hs_ints.intersections, hs_ints.dual_facets)
            # Only inlcude the facets corresponding to entries, not the boundaries
            total_entries = len(defects)
            ints_and_facets = filter(lambda int_and_facet: all(np.array(int_and_facet[1]) < total_entries),
                                     ints_and_facets)
            # sort based on transition level
            ints_and_facets = list(sorted(ints_and_facets, key=lambda int_and_facet: int_and_facet[0][0]))
            # Unpack into lists
            intersections, facets = zip(*ints_and_facets)
            # Map of transition level: charge states

            transition_level_map[defects[0].name] = {
                intersection[0]: [defects[i].charge for i in facet]
                for intersection, facet in ints_and_facets
            }

            stable_entries[defects[0].name] = [defects[i] for dual in facets for i in dual]

            finished_charges[defects[0].name] = [defect.charge for defect in defects]

        self.transition_level_map = transition_level_map
        self.transition_levels = {
            defect_name: list(defect_tls.keys())
            for defect_name, defect_tls in transition_level_map.items()
        }
        self.stable_entries = stable_entries
        self.finished_charges = finished_charges
        self.stable_charges = {
            defect_name: [entry.charge for entry in entries]
            for defect_name, entries in stable_entries.items()
        }

    def OLDfind_stable_charges(self):
        """
        Set the DefectPhaseDiagram attributes (which dont depend on chemical potential)

        THIS is a wrapper since find_stable_charges does not currently exist...will switch over when able
        """
        def similar_defect(a, b):
            """
            Used to filter out similar defects of different charges which
            are defined by the same type and location
            """
            return a.name == b.name and a.site == b.site

        self.stable_charges = {} #keys are defect names, items are list of charge states that are stable
        self.finished_charges = {} #keys are defect names, items are list of charge states that are included in the phase diagram
        self.transition_levels = {} #keys are defect names, items are list of [fermi level for transition, previous q, next q] sets
        xlim = (-0.1, self.band_gap+.1)
        nb_steps = 20 #SUPER bare bones because it is so slow...
        x = np.arange(xlim[0], xlim[1], (xlim[1]-xlim[0])/nb_steps)
        list_elts = [elt  for dfct in self.entries  for elt in dfct.defect.generate_defect_structure().composition.elements]
        no_chempots = {elt: 0. for elt in set(list_elts)} #zerod chemical potentials for calculating stable defects

        # Grouping by defect types
        for _, defects in groupby(sorted(self.entries, key=similar_defect), similar_defect):
            defects = list(defects)
            t = defects[0].name
            print('Analyzing {}'.format(t))
            trans_level = []
            chg_type = []
            prev_min_q, cur_min_q = None, None
            for x_step in x:
                miny = 10000
                for dfct in defects:
                    val = dfct.formation_energy(chemical_potentials = no_chempots, fermi_level=x_step)
                    if val < miny:
                        miny = val
                        cur_min_q = dfct.charge

                if prev_min_q is not None:
                    if cur_min_q != prev_min_q:
                        trans_level.append((x_step, prev_min_q, cur_min_q))
                    if cur_min_q not in chg_type:
                        chg_type.append(cur_min_q)
                prev_min_q = cur_min_q

            self.stable_charges[t] = chg_type[:]
            self.finished_charges[t] = [e.charge for e in self.entries if e.name == t]
            self.transition_levels[t] = trans_level[:]

    @property
    def defect_types(self):
        """
        List types of defects existing in the DefectPhaseDiagram
        """
        return list(self.finished_charges.keys())

    @property
    def all_stable_entries(self):
        """
        List all stable entries (defect+charge) in the DefectPhaseDiagram
        """
        return set(chain.from_iterable(self.stable_entries.values()))

    @property
    def all_unstable_entries(self):
        """
        List all unstable entries (defect+charge) in the DefectPhaseDiagram
        """
        all_stable_entries = self.all_stable_entries
        return [e for e in self.entries if e not in all_stable_entries]

    def defect_concentrations(self, chemical_potentials, temperature=300, fermi_level=0.):
        """
        Give list of all concentrations at specified efermi in the DefectPhaseDiagram
        args:
            chemical_potentials = {Element: number} is dictionary of chemical potentials to provide formation energies for
            temperature = temperature to produce concentrations from
            fermi_level: (float) is fermi level relative to valence band maximum
                Default efermi = 0 = VBM energy
        returns:
            list of dictionaries of defect concentrations
        """
        concentrations = []
        for dfct in self.all_stable_entries:
            concentrations.append({
                'conc':
                dfct.defect_concentration(
                    chemical_potentials=chemical_potentials, temperature=temperature, fermi_level=fermi_level),
                'name':
                dfct.name,
                'charge':
                dfct.charge
            })

            # TODO: Should be this entry: concentration?
        return concentrations

    def suggest_charges(self, tolerance=0.1):
        """
        Suggest possible charges for defects to computee based on proximity
        of known transitions from entires to VBM and CBM

        Args:
            tolerance (float): tolerance with respect to the VBM and CBM to
                    `          continue to compute new charges
        """
        recommendations = {}

        for def_type in self.defect_types:
            test_charges = np.arange(
                np.min(self.stable_charges[def_type]) - 1,
                np.max(self.stable_charges[def_type]) + 2)
            test_charges = [charge for charge in test_charges if charge not in self.finished_charges[def_type]]

            # More positive charges will shift the minimum transition level down
            # Max charge is limited by this if its transition level is close to VBM
            min_tl = min(self.transition_level_map[def_type].keys())
            if min_tl < tolerance:
                max_charge = max(self.transition_level_map[def_type][min_tl])
                test_charges = [charge for charge in test_charges if charge < max_charge]

            # More negative charges will shift the maximum transition level up
            # Minimum charge is limited by this if transition level is near CBM
            max_tl = max(self.transition_level_map[def_type].keys())
            if max_tl > (self.band_gap - tolerance):
                min_charge = min(self.transition_level_map[def_type][max_tl])
                test_charges = [charge for charge in test_charges if charge > min_charge]

            recommendations[def_type] = test_charges

        return recommendations

    def solve_for_fermi_energy(self, temperature, chemical_potentials, bulk_dos):
        """
        Solve for the Fermi energy self-consistently as a function of T
        and p_O2
        Observations are Defect concentrations, electron and hole conc
        Args:
            bulk_dos: bulk system dos (pymatgen Dos object)
            gap: Can be used to specify experimental gap.
                Will be useful if the self consistent Fermi level
                is > DFT gap
        Returns:
            Fermi energy
        """

        fdos = FermiDos(bulk_dos, bandgap=self.band_gap)

        def _get_total_q(ef):

            qd_tot = sum([
                d['charge'] * d['conc']
                for d in self.list_defect_concentrations(
                    chemical_potentials=chemical_potentials, temperature=temperature, fermi_level=ef)
            ])
            qd_tot += fdos.get_doping(fermi=ef, T=temperature)
            return qd_tot

        return bisect(_get_total_q, -1., self.band_gap + 1.)

    def solve_non_eq_fermilevel(self, lowT, highT, bulk_dos, show_carriers=True, hold_htemp_ef=True):
        """
        Solve for the Fermi energy at a low temperature value, but with defect concentrations frozen in
            at some higher temperature value This means that only the fermi level and free charge carriers concentrations
            will be solved self consistently (simulating a 'frozen-in' defect concentration approach)

        Args:
            lowT: the temperature at which you want to solve for fermi level and free carrier concentrations
            highT: the fixed temperature that sets the defect concentrations at a higher temperature
            bulk_dos: bulk system dos (pymatgen Dos object)
            show_carriers: Dictates the outputs of this function
                Set to True if you want the new carrier concentrations in addition to the fermi levels
            hold_htemp_ef: Specifies whether the fermi level will be fixed for the defect concentrations at high T
                Set to False if you want the fermi level of the high temperature defects to be the same as the fermi level
                for the low temperature charge carriers (see below note on implementation for explanation of this flag)
        Returns:
            Fermi energy

        A NOTE ON IMPLEMENTATION: while testing this function, it became apparent that negative defect formation energies
            can sometimes cause stochastic results for the fermi level as a result of the defect concentrations not
            being allowed to change within the self consistent result. To circumvent this, this function has a flag ('hold_htemp_ef')
            to allow for the fermi level of high temperature defects to vary within the self consistent approach.
            While this is no longer the fully physical 'frozen-in' approach, it helps remove the stocastic results
            that result from numerical issues

        """
        from scipy.optimize import bisect
        q_h_cont = IntrinsicCarrier(bulk_dos, exp_gap=self.band_gap)

        if hold_htemp_ef:
            self.solve_for_fermi_energy(bulk_dos, gap=self.band_gap)
            defectlist = self.get_defects_concentration(highT, self._ef, unitcell=False)

        def noneq_total_q(ef):
            qd_tot = 0
            if hold_htemp_ef:
                for d in defectlist:
                    qd_tot += d['charge'] * d['conc']
            else:
                for d in self.get_defects_concentration(highT, ef, unitcell=False):
                    qd_tot += d['charge'] * d['conc']
            ef_ref_cbm = ef - self.band_gap
            nden = q_h_cont.get_n_density(ef_ref_cbm, lowT, ref='CBM', unitcell=False)
            pden = q_h_cont.get_p_density(ef, lowT, unitcell=False)
            qd_tot += pden - nden
            return qd_tot

        non_eq_ef = bisect(lambda e: noneq_total_q(e), -1., self.band_gap + 1.)
        if show_carriers:
            ef_ref_cbm = non_eq_ef - self.band_gap
            nden = q_h_cont.get_n_density(ef_ref_cbm, lowT, ref='CBM', unitcell=False)
            pden = q_h_cont.get_p_density(non_eq_ef, lowT, unitcell=False)
            return non_eq_ef, nden, pden
        else:
            return non_eq_ef

<|endoftext|>


<commit_before>from OpenGL.GL import *
from OpenGL.arrays import vbo
from .. GLGraphicsItem import GLGraphicsItem
from .. import shaders
from ... import QtGui
import numpy as np

__all__ = ['GLScatterPlotItem']

class GLScatterPlotItem(GLGraphicsItem):
    """Draws points at a list of 3D positions."""

    def __init__(self, **kwds):
        GLGraphicsItem.__init__(self)
        glopts = kwds.pop('glOptions', 'additive')
        self.setGLOptions(glopts)
        self.pos = []
        self.size = 10
        self.color = [1.0,1.0,1.0,0.5]
        self.pxMode = True
        #self.vbo = {}      ## VBO does not appear to improve performance very much.
        self.setData(**kwds)

    def setData(self, **kwds):
        """
        Update the data displayed by this item. All arguments are optional;
        for example it is allowed to update spot positions while leaving
        colors unchanged, etc.

        ====================  ==================================================
        **Arguments:**
        pos                   (N,3) array of floats specifying point locations.
        color                 (N,4) array of floats (0.0-1.0) specifying
                              spot colors OR a tuple of floats specifying
                              a single color for all spots.
        size                  (N,) array of floats specifying spot sizes or
                              a single value to apply to all spots.
        pxMode                If True, spot sizes are expressed in pixels.
                              Otherwise, they are expressed in item coordinates.
        ====================  ==================================================
        """
        args = ['pos', 'color', 'size', 'pxMode']
        for k in kwds.keys():
            if k not in args:
                raise Exception('Invalid keyword argument: %s (allowed arguments are %s)' % (k, str(args)))

        args.remove('pxMode')
        for arg in args:
            if arg in kwds:
                setattr(self, arg, kwds[arg])
                #self.vbo.pop(arg, None)

        self.pxMode = kwds.get('pxMode', self.pxMode)
        self.update()

    def initializeGL(self):

        ## Generate texture for rendering points
        w = 64
        def fn(x,y):
            r = ((x-w/2.)**2 + (y-w/2.)**2) ** 0.5
            return 255 * (w/2. - np.clip(r, w/2.-1.0, w/2.))
        pData = np.empty((w, w, 4))
        pData[:] = 255
        pData[:,:,3] = np.fromfunction(fn, pData.shape[:2])
        #print pData.shape, pData.min(), pData.max()
        pData = pData.astype(np.ubyte)

        self.pointTexture = glGenTextures(1)
        glActiveTexture(GL_TEXTURE0)
        glEnable(GL_TEXTURE_2D)
        glBindTexture(GL_TEXTURE_2D, self.pointTexture)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, pData.shape[0], pData.shape[1], 0, GL_RGBA, GL_UNSIGNED_BYTE, pData)

        self.shader = shaders.getShaderProgram('pointSprite')

    #def getVBO(self, name):
        #if name not in self.vbo:
            #self.vbo[name] = vbo.VBO(getattr(self, name).astype('f'))
        #return self.vbo[name]

    #def setupGLState(self):
        #"""Prepare OpenGL state for drawing. This function is called immediately before painting."""
        ##glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)  ## requires z-sorting to render properly.
        #glBlendFunc(GL_SRC_ALPHA, GL_ONE)
        #glEnable( GL_BLEND )
        #glEnable( GL_ALPHA_TEST )
        #glDisable( GL_DEPTH_TEST )

        ##glEnable( GL_POINT_SMOOTH )

        ##glHint(GL_POINT_SMOOTH_HINT, GL_NICEST)
        ##glPointParameterfv(GL_POINT_DISTANCE_ATTENUATION, (0, 0, -1e-3))
        ##glPointParameterfv(GL_POINT_SIZE_MAX, (65500,))
        ##glPointParameterfv(GL_POINT_SIZE_MIN, (0,))

    def paint(self):
        self.setupGLState()

        glEnable(GL_POINT_SPRITE)

        glActiveTexture(GL_TEXTURE0)
        glEnable( GL_TEXTURE_2D )
        glBindTexture(GL_TEXTURE_2D, self.pointTexture)

        glTexEnvi(GL_POINT_SPRITE, GL_COORD_REPLACE, GL_TRUE)
        #glTexEnvi(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_REPLACE)    ## use texture color exactly
        #glTexEnvf( GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_MODULATE )  ## texture modulates current color
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
        glEnable(GL_PROGRAM_POINT_SIZE)

        with self.shader:
            #glUniform1i(self.shader.uniform('texture'), 0)  ## inform the shader which texture to use
            glEnableClientState(GL_VERTEX_ARRAY)
            try:
                pos = self.pos
                #if pos.ndim > 2:
                    #pos = pos.reshape((reduce(lambda a,b: a*b, pos.shape[:-1]), pos.shape[-1]))
                glVertexPointerf(pos)

                if isinstance(self.color, np.ndarray):
                    glEnableClientState(GL_COLOR_ARRAY)
                    glColorPointerf(self.color)
                else:
                    if isinstance(self.color, QtGui.QColor):
                        glColor4f(*fn.glColor(self.color))
                    else:
                        glColor4f(*self.color)

                if not self.pxMode or isinstance(self.size, np.ndarray):
                    glEnableClientState(GL_NORMAL_ARRAY)
                    norm = np.empty(pos.shape)
                    if self.pxMode:
                        norm[...,0] = self.size
                    else:
                        gpos = self.mapToView(pos.transpose()).transpose()
                        pxSize = self.view().pixelSize(gpos)
                        norm[...,0] = self.size / pxSize

                    glNormalPointerf(norm)
                else:
                    glNormal3f(self.size, 0, 0)  ## vertex shader uses norm.x to determine point size
                    #glPointSize(self.size)
                glDrawArrays(GL_POINTS, 0, int(pos.size / pos.shape[-1]))
            finally:
                glDisableClientState(GL_NORMAL_ARRAY)
                glDisableClientState(GL_VERTEX_ARRAY)
                glDisableClientState(GL_COLOR_ARRAY)
                #posVBO.unbind()

        #for i in range(len(self.pos)):
            #pos = self.pos[i]

            #if isinstance(self.color, np.ndarray):
                #color = self.color[i]
            #else:
                #color = self.color
            #if isinstance(self.color, QtGui.QColor):
                #color = fn.glColor(self.color)

            #if isinstance(self.size, np.ndarray):
                #size = self.size[i]
            #else:
                #size = self.size

            #pxSize = self.view().pixelSize(QtGui.QVector3D(*pos))

            #glPointSize(size / pxSize)
            #glBegin( GL_POINTS )
            #glColor4f(*color)  # x is blue
            ##glNormal3f(size, 0, 0)
            #glVertex3f(*pos)
            #glEnd()

<commit_msg>Fix memory leak in GLScatterPlotItem<commit_after>from OpenGL.GL import *
from OpenGL.arrays import vbo
from .. GLGraphicsItem import GLGraphicsItem
from .. import shaders
from ... import QtGui
import numpy as np

__all__ = ['GLScatterPlotItem']

class GLScatterPlotItem(GLGraphicsItem):
    """Draws points at a list of 3D positions."""

    def __init__(self, **kwds):
        GLGraphicsItem.__init__(self)
        glopts = kwds.pop('glOptions', 'additive')
        self.setGLOptions(glopts)
        self.pos = []
        self.size = 10
        self.color = [1.0,1.0,1.0,0.5]
        self.pxMode = True
        #self.vbo = {}      ## VBO does not appear to improve performance very much.
        self.setData(**kwds)

    def setData(self, **kwds):
        """
        Update the data displayed by this item. All arguments are optional;
        for example it is allowed to update spot positions while leaving
        colors unchanged, etc.

        ====================  ==================================================
        **Arguments:**
        pos                   (N,3) array of floats specifying point locations.
        color                 (N,4) array of floats (0.0-1.0) specifying
                              spot colors OR a tuple of floats specifying
                              a single color for all spots.
        size                  (N,) array of floats specifying spot sizes or
                              a single value to apply to all spots.
        pxMode                If True, spot sizes are expressed in pixels.
                              Otherwise, they are expressed in item coordinates.
        ====================  ==================================================
        """
        args = ['pos', 'color', 'size', 'pxMode']
        for k in kwds.keys():
            if k not in args:
                raise Exception('Invalid keyword argument: %s (allowed arguments are %s)' % (k, str(args)))

        args.remove('pxMode')
        for arg in args:
            if arg in kwds:
                setattr(self, arg, kwds[arg])
                #self.vbo.pop(arg, None)

        self.pxMode = kwds.get('pxMode', self.pxMode)
        self.update()

    def initializeGL(self):

        ## Generate texture for rendering points
        w = 64
        def fn(x,y):
            r = ((x-w/2.)**2 + (y-w/2.)**2) ** 0.5
            return 255 * (w/2. - np.clip(r, w/2.-1.0, w/2.))
        pData = np.empty((w, w, 4))
        pData[:] = 255
        pData[:,:,3] = np.fromfunction(fn, pData.shape[:2])
        #print pData.shape, pData.min(), pData.max()
        pData = pData.astype(np.ubyte)

        if getattr(self, "pointTexture", None) is None:
            self.pointTexture = glGenTextures(1)
        glActiveTexture(GL_TEXTURE0)
        glEnable(GL_TEXTURE_2D)
        glBindTexture(GL_TEXTURE_2D, self.pointTexture)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, pData.shape[0], pData.shape[1], 0, GL_RGBA, GL_UNSIGNED_BYTE, pData)

        self.shader = shaders.getShaderProgram('pointSprite')

    #def getVBO(self, name):
        #if name not in self.vbo:
            #self.vbo[name] = vbo.VBO(getattr(self, name).astype('f'))
        #return self.vbo[name]

    #def setupGLState(self):
        #"""Prepare OpenGL state for drawing. This function is called immediately before painting."""
        ##glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)  ## requires z-sorting to render properly.
        #glBlendFunc(GL_SRC_ALPHA, GL_ONE)
        #glEnable( GL_BLEND )
        #glEnable( GL_ALPHA_TEST )
        #glDisable( GL_DEPTH_TEST )

        ##glEnable( GL_POINT_SMOOTH )

        ##glHint(GL_POINT_SMOOTH_HINT, GL_NICEST)
        ##glPointParameterfv(GL_POINT_DISTANCE_ATTENUATION, (0, 0, -1e-3))
        ##glPointParameterfv(GL_POINT_SIZE_MAX, (65500,))
        ##glPointParameterfv(GL_POINT_SIZE_MIN, (0,))

    def paint(self):
        self.setupGLState()

        glEnable(GL_POINT_SPRITE)

        glActiveTexture(GL_TEXTURE0)
        glEnable( GL_TEXTURE_2D )
        glBindTexture(GL_TEXTURE_2D, self.pointTexture)

        glTexEnvi(GL_POINT_SPRITE, GL_COORD_REPLACE, GL_TRUE)
        #glTexEnvi(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_REPLACE)    ## use texture color exactly
        #glTexEnvf( GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_MODULATE )  ## texture modulates current color
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
        glEnable(GL_PROGRAM_POINT_SIZE)

        with self.shader:
            #glUniform1i(self.shader.uniform('texture'), 0)  ## inform the shader which texture to use
            glEnableClientState(GL_VERTEX_ARRAY)
            try:
                pos = self.pos
                #if pos.ndim > 2:
                    #pos = pos.reshape((reduce(lambda a,b: a*b, pos.shape[:-1]), pos.shape[-1]))
                glVertexPointerf(pos)

                if isinstance(self.color, np.ndarray):
                    glEnableClientState(GL_COLOR_ARRAY)
                    glColorPointerf(self.color)
                else:
                    if isinstance(self.color, QtGui.QColor):
                        glColor4f(*fn.glColor(self.color))
                    else:
                        glColor4f(*self.color)

                if not self.pxMode or isinstance(self.size, np.ndarray):
                    glEnableClientState(GL_NORMAL_ARRAY)
                    norm = np.empty(pos.shape)
                    if self.pxMode:
                        norm[...,0] = self.size
                    else:
                        gpos = self.mapToView(pos.transpose()).transpose()
                        pxSize = self.view().pixelSize(gpos)
                        norm[...,0] = self.size / pxSize

                    glNormalPointerf(norm)
                else:
                    glNormal3f(self.size, 0, 0)  ## vertex shader uses norm.x to determine point size
                    #glPointSize(self.size)
                glDrawArrays(GL_POINTS, 0, int(pos.size / pos.shape[-1]))
            finally:
                glDisableClientState(GL_NORMAL_ARRAY)
                glDisableClientState(GL_VERTEX_ARRAY)
                glDisableClientState(GL_COLOR_ARRAY)
                #posVBO.unbind()

        #for i in range(len(self.pos)):
            #pos = self.pos[i]

            #if isinstance(self.color, np.ndarray):
                #color = self.color[i]
            #else:
                #color = self.color
            #if isinstance(self.color, QtGui.QColor):
                #color = fn.glColor(self.color)

            #if isinstance(self.size, np.ndarray):
                #size = self.size[i]
            #else:
                #size = self.size

            #pxSize = self.view().pixelSize(QtGui.QVector3D(*pos))

            #glPointSize(size / pxSize)
            #glBegin( GL_POINTS )
            #glColor4f(*color)  # x is blue
            ##glNormal3f(size, 0, 0)
            #glVertex3f(*pos)
            #glEnd()

<|endoftext|>


<commit_before>import pytest

import pyramid.testing
from pyramid.authorization import ACLAuthorizationPolicy
from pyramid.interfaces import IAuthenticationPolicy

from zope.interface import implementer

from zope.interface.verify import (
        verifyClass,
        verifyObject
        )

from pyramid_authsanity.interfaces import (
        IAuthSourceService,
        IAuthService,
        )

from pyramid_authsanity.policy import NoAuthCompleted

def test_clean_principal_invalid():
    from pyramid_authsanity.policy import _clean_principal
    from pyramid.security import Everyone

    ret = _clean_principal(Everyone)

    assert ret == None

def test_clean_principal_valid():
    from pyramid_authsanity.policy import _clean_principal

    ret = _clean_principal('root')

    assert ret == 'root'

class TestAuthServicePolicyInterface(object):
    def test_verify(self):
        from pyramid_authsanity.policy import AuthServicePolicy
        assert verifyClass(IAuthenticationPolicy, AuthServicePolicy)
        assert verifyObject(IAuthenticationPolicy, AuthServicePolicy())

class TestAuthServicePolicy(object):
    @pytest.fixture(autouse=True)
    def pyramid_config(self, request):
        from pyramid.interfaces import IDebugLogger
        self.config = pyramid.testing.setUp()
        self.config.set_authorization_policy(ACLAuthorizationPolicy())
        self.logger = DummyLogger()
        self.config.registry.registerUtility(self.logger, IDebugLogger)

        def finish():
            del self.config
            pyramid.testing.tearDown()
        request.addfinalizer(finish)

    def _makeOne(self, debug=False, source=None, auth=None):
        from pyramid_authsanity import AuthServicePolicy
        def find_services(request):
            return (source, auth)

        def session_registered(request):
            return False

        policy = AuthServicePolicy(debug=debug)
        policy._find_services = find_services
        policy._session_registered = session_registered
        return policy

    def _makeOneRequest(self):
        request = DummyRequest()
        request.registry = self.config.registry
        return request

    def test_find_services(self):
        policy = self._makeOne()
        request = self._makeOneRequest()

        (source, auth) = policy._find_services(request)

        assert source == None
        assert auth == None

    def test_fake_source_ticket(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(None)(context, request)
        auth = fake_auth_init()(context, request)

        assert verifyObject(IAuthService, auth)
        assert verifyObject(IAuthSourceService, source)

    def test_valid_source_ticket(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])(context, request)
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])(context, request)

        policy = self._makeOne(debug=True, source=source, auth=auth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == 'test'

    def test_invalid_source_ticket(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'invalid'])(context, request)
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])(context, request)

        policy = self._makeOne(source=source, auth=auth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == None

    def test_bad_auth_service(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])(context, request)

        class BadAuth(object):
            def userid(self):
                raise ValueError('No ticket verified')

            def verify_ticket(self, principal, ticket):
                pass

        auth = BadAuth()

        policy = self._makeOne(source=source, auth=auth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == None

    def test_no_user_effective_principals(self):
        from pyramid.security import Everyone
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])(context, request)
        auth = fake_auth_init()(context, request)

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone] == groups

    def test_user_bad_principal_effective_principals(self):
        from pyramid.security import Everyone
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([Everyone, 'valid'])(context, request)
        auth = fake_auth_init(fake_userid=Everyone, valid_tickets=['valid'])(context, request)

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone] == groups

    def test_effective_principals(self):
        from pyramid.security import Everyone, Authenticated
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])(context, request)
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])(context, request)

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone, Authenticated, 'test'] == groups

    def test_remember(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])(context, request)
        auth = fake_auth_init(fake_userid='test')(context, request)

        policy = self._makeOne(source=source, auth=auth)

        headers = policy.remember(request, 'test')

        assert len(headers) == 0
        assert len(auth.valid_tickets) >= 1

    def test_forget(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])(context, request)
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])(context, request)

        policy = self._makeOne(source=source, auth=auth)

        assert 'valid' in auth.valid_tickets

        headers = policy.forget(request)

        assert len(headers) == 0
        assert 'valid' not in auth.valid_tickets

class TestAuthServicePolicyIntegration(object):
    @pytest.fixture(autouse=True)
    def pyramid_config(self, request):
        from pyramid.interfaces import IDebugLogger, ISessionFactory
        self.config = pyramid.testing.setUp()
        self.config.registry.registerUtility(lambda: None, ISessionFactory)
        self.config.include('pyramid_services')
        self.config.set_authorization_policy(ACLAuthorizationPolicy())
        self.logger = DummyLogger()
        self.config.registry.registerUtility(self.logger, IDebugLogger)

        def finish():
            del self.config
            pyramid.testing.tearDown()
        request.addfinalizer(finish)

    def _makeOne(self, debug=False, source=None, auth=None):
        from pyramid_authsanity import AuthServicePolicy

        if source:
            self.config.register_service_factory(source, iface=IAuthSourceService)

        if auth:
            self.config.register_service_factory(auth, iface=IAuthService)

        return AuthServicePolicy(debug=debug)

    def _makeOneRequest(self):
        from pyramid_services import find_service
        from zope.interface.adapter import AdapterRegistry
        import types

        request = DummyRequest()
        request.find_service = types.MethodType(find_service, request)
        request.registry = self.config.registry
        request.service_cache = AdapterRegistry()

        return request

    def test_include_me(self):
        from pyramid_authsanity.policy import AuthServicePolicy
        self.config.include('pyramid_authsanity')
        self.config.commit()
        introspector = self.config.registry.introspector
        auth_policy = introspector.get('authentication policy', None)

        assert isinstance(auth_policy['policy'], AuthServicePolicy)

    def test_logging(self):
        policy = self._makeOne(debug=True)
        request = DummyRequest(registry=self.config.registry)
        policy._log('this message', 'test_logging', request)

        assert len(self.logger.logentries) >= 1

    def test_find_services(self):
        from pyramid_authsanity.interfaces import (IAuthSourceService, IAuthService)
        self.config.register_service_factory(lambda x,y: 'Source', iface=IAuthSourceService)
        self.config.register_service_factory(lambda x,y: 'Auth', iface=IAuthService)

        policy = self._makeOne()
        request = self._makeOneRequest()

        (sourcesvc, authsvc) = policy._find_services(request)

        assert sourcesvc == 'Source'
        assert authsvc == 'Auth'

    def test_valid_source_ticket(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])

        policy = self._makeOne(debug=True, source=source, auth=auth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == 'test'

    def test_invalid_source_ticket(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'invalid'])
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])

        policy = self._makeOne(source=source, auth=auth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == None

    def test_bad_auth_service(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])

        class BadAuth(object):
            def __init__(self, context, request):
                pass

            def userid(self):
                raise ValueError('No ticket verification')

            def verify_ticket(self, principal, ticket):
                pass

        policy = self._makeOne(source=source, auth=BadAuth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == None

    def test_no_user_effective_principals(self):
        from pyramid.security import Everyone
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])
        auth = fake_auth_init()

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone] == groups

    def test_user_bad_principal_effective_principals(self):
        from pyramid.security import Everyone
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([Everyone, 'valid'])
        auth = fake_auth_init(fake_userid=Everyone, valid_tickets=['valid'])

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone] == groups

    def test_effective_principals(self):
        from pyramid.security import Everyone, Authenticated
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone, Authenticated, 'test'] == groups

    def test_remember(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])
        auth = fake_auth_init(fake_userid='test')

        policy = self._makeOne(source=source, auth=auth)

        headers = policy.remember(request, 'test')

        authreq = request.find_service(IAuthService)

        assert len(headers) == 0
        assert len(authreq.valid_tickets) >= 1

    def test_remember_value_json_serializable(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])
        auth = fake_auth_init(fake_userid='test')

        policy = self._makeOne(source=source, auth=auth)

        headers = policy.remember(request, 'test')

        authreq = request.find_service(IAuthService)
        sourcereq = request.find_service(IAuthSourceService)

        assert len(headers) == 0
        assert len(authreq.valid_tickets) >= 1
        import json

        assert json.dumps(sourcereq.value)

    def test_remember_same_user(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid_ticket'])
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid_ticket'])

        policy = self._makeOne(source=source, auth=auth)

        headers = policy.remember(request, 'test')

        authreq = request.find_service(IAuthService)

        assert len(headers) == 0
        assert len(authreq.valid_tickets) >= 1

    def test_forget(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])

        policy = self._makeOne(source=source, auth=auth)

        authreq = request.find_service(IAuthService)

        assert 'valid' in authreq.valid_tickets

        headers = policy.forget(request)

        assert len(headers) == 0
        assert 'valid' not in authreq.valid_tickets

class DummyLogger(object):
    def debug(self, log):
        self.logentries.append(log)

    def __init__(self):
        self.logentries = []

class DummyRequest(object):
    domain = 'example.net'

    def __init__(self, environ=None, session=None, registry=None, cookie=None):
        class Session(dict):
            def invalidate(self):
                self.clear()

            def new_csrf_token(self):
                pass

        self.environ = environ or {}
        self.session = Session()
        self.session.update(session or {})
        self.registry = registry
        self.callbacks = []
        self.cookies = cookie or []
        self.context = None

    def add_response_callback(self, callback):
        self.callbacks.append(callback)

def fake_source_init(fake_value):
    @implementer(IAuthSourceService)
    class fake_source(object):
        vary = ['Cookie']
        def __init__(self, context, request):
            self.value = fake_value

        def get_value(self):
            return self.value if self.value else [None, None]

        def headers_remember(self, value):
            self.value = value
            return []

        def headers_forget(self):
            self.value = [None, None]
            return []

    return fake_source

def fake_auth_init(fake_userid=None, fake_groups=list(), valid_tickets=list()):
    @implementer(IAuthService)
    class fake_auth(object):

        def __init__(self, context, request):
            self.authcomplete = False
            self.ticketvalid = False
            self._userid = fake_userid
            self._groups = fake_groups
            self.valid_tickets = valid_tickets

        def userid(self):
            if not self.authcomplete:
                raise ValueError('No ticket verified.')

            if not self.ticketvalid:
                return None

            return self._userid

        def groups(self):
            return self._groups

        def verify_ticket(self, principal, ticket):
            self.authcomplete = True
            if principal == self._userid and ticket in self.valid_tickets:
                self.ticketvalid = True

        def add_ticket(self, principal, ticket):
            if ticket not in self.valid_tickets:
                self.valid_tickets.append(ticket)

        def remove_ticket(self, ticket):
            if ticket in self.valid_tickets:
                self.valid_tickets.remove(ticket)

    return fake_auth
<commit_msg>Source needs to store two values, currently in a list<commit_after>import pytest

import pyramid.testing
from pyramid.authorization import ACLAuthorizationPolicy
from pyramid.interfaces import IAuthenticationPolicy

from zope.interface import implementer

from zope.interface.verify import (
        verifyClass,
        verifyObject
        )

from pyramid_authsanity.interfaces import (
        IAuthSourceService,
        IAuthService,
        )

from pyramid_authsanity.policy import NoAuthCompleted

def test_clean_principal_invalid():
    from pyramid_authsanity.policy import _clean_principal
    from pyramid.security import Everyone

    ret = _clean_principal(Everyone)

    assert ret == None

def test_clean_principal_valid():
    from pyramid_authsanity.policy import _clean_principal

    ret = _clean_principal('root')

    assert ret == 'root'

class TestAuthServicePolicyInterface(object):
    def test_verify(self):
        from pyramid_authsanity.policy import AuthServicePolicy
        assert verifyClass(IAuthenticationPolicy, AuthServicePolicy)
        assert verifyObject(IAuthenticationPolicy, AuthServicePolicy())

class TestAuthServicePolicy(object):
    @pytest.fixture(autouse=True)
    def pyramid_config(self, request):
        from pyramid.interfaces import IDebugLogger
        self.config = pyramid.testing.setUp()
        self.config.set_authorization_policy(ACLAuthorizationPolicy())
        self.logger = DummyLogger()
        self.config.registry.registerUtility(self.logger, IDebugLogger)

        def finish():
            del self.config
            pyramid.testing.tearDown()
        request.addfinalizer(finish)

    def _makeOne(self, debug=False, source=None, auth=None):
        from pyramid_authsanity import AuthServicePolicy
        def find_services(request):
            return (source, auth)

        def session_registered(request):
            return False

        policy = AuthServicePolicy(debug=debug)
        policy._find_services = find_services
        policy._session_registered = session_registered
        return policy

    def _makeOneRequest(self):
        request = DummyRequest()
        request.registry = self.config.registry
        return request

    def test_find_services(self):
        policy = self._makeOne()
        request = self._makeOneRequest()

        (source, auth) = policy._find_services(request)

        assert source == None
        assert auth == None

    def test_fake_source_ticket(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(None)(context, request)
        auth = fake_auth_init()(context, request)

        assert verifyObject(IAuthService, auth)
        assert verifyObject(IAuthSourceService, source)

    def test_valid_source_ticket(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])(context, request)
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])(context, request)

        policy = self._makeOne(debug=True, source=source, auth=auth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == 'test'

    def test_invalid_source_ticket(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'invalid'])(context, request)
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])(context, request)

        policy = self._makeOne(source=source, auth=auth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == None

    def test_bad_auth_service(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])(context, request)

        class BadAuth(object):
            def userid(self):
                raise ValueError('No ticket verified')

            def verify_ticket(self, principal, ticket):
                pass

        auth = BadAuth()

        policy = self._makeOne(source=source, auth=auth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == None

    def test_no_user_effective_principals(self):
        from pyramid.security import Everyone
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])(context, request)
        auth = fake_auth_init()(context, request)

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone] == groups

    def test_user_bad_principal_effective_principals(self):
        from pyramid.security import Everyone
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([Everyone, 'valid'])(context, request)
        auth = fake_auth_init(fake_userid=Everyone, valid_tickets=['valid'])(context, request)

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone] == groups

    def test_effective_principals(self):
        from pyramid.security import Everyone, Authenticated
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])(context, request)
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])(context, request)

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone, Authenticated, 'test'] == groups

    def test_remember(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])(context, request)
        auth = fake_auth_init(fake_userid='test')(context, request)

        policy = self._makeOne(source=source, auth=auth)

        headers = policy.remember(request, 'test')

        assert len(headers) == 0
        assert len(auth.valid_tickets) >= 1
        assert isinstance(source.value, list)
        assert len(source.value) == 2

    def test_forget(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])(context, request)
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])(context, request)

        policy = self._makeOne(source=source, auth=auth)

        assert 'valid' in auth.valid_tickets

        headers = policy.forget(request)

        assert len(headers) == 0
        assert 'valid' not in auth.valid_tickets

class TestAuthServicePolicyIntegration(object):
    @pytest.fixture(autouse=True)
    def pyramid_config(self, request):
        from pyramid.interfaces import IDebugLogger, ISessionFactory
        self.config = pyramid.testing.setUp()
        self.config.registry.registerUtility(lambda: None, ISessionFactory)
        self.config.include('pyramid_services')
        self.config.set_authorization_policy(ACLAuthorizationPolicy())
        self.logger = DummyLogger()
        self.config.registry.registerUtility(self.logger, IDebugLogger)

        def finish():
            del self.config
            pyramid.testing.tearDown()
        request.addfinalizer(finish)

    def _makeOne(self, debug=False, source=None, auth=None):
        from pyramid_authsanity import AuthServicePolicy

        if source:
            self.config.register_service_factory(source, iface=IAuthSourceService)

        if auth:
            self.config.register_service_factory(auth, iface=IAuthService)

        return AuthServicePolicy(debug=debug)

    def _makeOneRequest(self):
        from pyramid_services import find_service
        from zope.interface.adapter import AdapterRegistry
        import types

        request = DummyRequest()
        request.find_service = types.MethodType(find_service, request)
        request.registry = self.config.registry
        request.service_cache = AdapterRegistry()

        return request

    def test_include_me(self):
        from pyramid_authsanity.policy import AuthServicePolicy
        self.config.include('pyramid_authsanity')
        self.config.commit()
        introspector = self.config.registry.introspector
        auth_policy = introspector.get('authentication policy', None)

        assert isinstance(auth_policy['policy'], AuthServicePolicy)

    def test_logging(self):
        policy = self._makeOne(debug=True)
        request = DummyRequest(registry=self.config.registry)
        policy._log('this message', 'test_logging', request)

        assert len(self.logger.logentries) >= 1

    def test_find_services(self):
        from pyramid_authsanity.interfaces import (IAuthSourceService, IAuthService)
        self.config.register_service_factory(lambda x,y: 'Source', iface=IAuthSourceService)
        self.config.register_service_factory(lambda x,y: 'Auth', iface=IAuthService)

        policy = self._makeOne()
        request = self._makeOneRequest()

        (sourcesvc, authsvc) = policy._find_services(request)

        assert sourcesvc == 'Source'
        assert authsvc == 'Auth'

    def test_valid_source_ticket(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])

        policy = self._makeOne(debug=True, source=source, auth=auth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == 'test'

    def test_invalid_source_ticket(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'invalid'])
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])

        policy = self._makeOne(source=source, auth=auth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == None

    def test_bad_auth_service(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])

        class BadAuth(object):
            def __init__(self, context, request):
                pass

            def userid(self):
                raise ValueError('No ticket verification')

            def verify_ticket(self, principal, ticket):
                pass

        policy = self._makeOne(source=source, auth=BadAuth)

        authuserid = policy.authenticated_userid(request)
        assert authuserid == None

    def test_no_user_effective_principals(self):
        from pyramid.security import Everyone
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])
        auth = fake_auth_init()

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone] == groups

    def test_user_bad_principal_effective_principals(self):
        from pyramid.security import Everyone
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([Everyone, 'valid'])
        auth = fake_auth_init(fake_userid=Everyone, valid_tickets=['valid'])

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone] == groups

    def test_effective_principals(self):
        from pyramid.security import Everyone, Authenticated
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])

        policy = self._makeOne(source=source, auth=auth)

        groups = policy.effective_principals(request)

        assert [Everyone, Authenticated, 'test'] == groups

    def test_remember(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])
        auth = fake_auth_init(fake_userid='test')

        policy = self._makeOne(source=source, auth=auth)

        headers = policy.remember(request, 'test')

        authreq = request.find_service(IAuthService)

        assert len(headers) == 0
        assert len(authreq.valid_tickets) >= 1

    def test_remember_value_json_serializable(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init([None, None])
        auth = fake_auth_init(fake_userid='test')

        policy = self._makeOne(source=source, auth=auth)

        headers = policy.remember(request, 'test')

        authreq = request.find_service(IAuthService)
        sourcereq = request.find_service(IAuthSourceService)

        assert len(headers) == 0
        assert len(authreq.valid_tickets) >= 1
        import json

        assert json.dumps(sourcereq.value)

    def test_remember_same_user(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid_ticket'])
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid_ticket'])

        policy = self._makeOne(source=source, auth=auth)

        headers = policy.remember(request, 'test')

        authreq = request.find_service(IAuthService)

        assert len(headers) == 0
        assert len(authreq.valid_tickets) >= 1

    def test_forget(self):
        context = None
        request = self._makeOneRequest()
        source = fake_source_init(['test', 'valid'])
        auth = fake_auth_init(fake_userid='test', valid_tickets=['valid'])

        policy = self._makeOne(source=source, auth=auth)

        authreq = request.find_service(IAuthService)

        assert 'valid' in authreq.valid_tickets

        headers = policy.forget(request)

        assert len(headers) == 0
        assert 'valid' not in authreq.valid_tickets

class DummyLogger(object):
    def debug(self, log):
        self.logentries.append(log)

    def __init__(self):
        self.logentries = []

class DummyRequest(object):
    domain = 'example.net'

    def __init__(self, environ=None, session=None, registry=None, cookie=None):
        class Session(dict):
            def invalidate(self):
                self.clear()

            def new_csrf_token(self):
                pass

        self.environ = environ or {}
        self.session = Session()
        self.session.update(session or {})
        self.registry = registry
        self.callbacks = []
        self.cookies = cookie or []
        self.context = None

    def add_response_callback(self, callback):
        self.callbacks.append(callback)

def fake_source_init(fake_value):
    @implementer(IAuthSourceService)
    class fake_source(object):
        vary = ['Cookie']
        def __init__(self, context, request):
            self.value = fake_value

        def get_value(self):
            return self.value if self.value else [None, None]

        def headers_remember(self, value):
            self.value = value
            return []

        def headers_forget(self):
            self.value = [None, None]
            return []

    return fake_source

def fake_auth_init(fake_userid=None, fake_groups=list(), valid_tickets=list()):
    @implementer(IAuthService)
    class fake_auth(object):

        def __init__(self, context, request):
            self.authcomplete = False
            self.ticketvalid = False
            self._userid = fake_userid
            self._groups = fake_groups
            self.valid_tickets = valid_tickets

        def userid(self):
            if not self.authcomplete:
                raise ValueError('No ticket verified.')

            if not self.ticketvalid:
                return None

            return self._userid

        def groups(self):
            return self._groups

        def verify_ticket(self, principal, ticket):
            self.authcomplete = True
            if principal == self._userid and ticket in self.valid_tickets:
                self.ticketvalid = True

        def add_ticket(self, principal, ticket):
            if ticket not in self.valid_tickets:
                self.valid_tickets.append(ticket)

        def remove_ticket(self, ticket):
            if ticket in self.valid_tickets:
                self.valid_tickets.remove(ticket)

    return fake_auth
<|endoftext|>


<commit_before>import ctypes
import re

def ValidHandle(value):
    if value == 0:
        raise ctypes.WinError()
    return value

import serial
from serial.win32 import ULONG_PTR
from ctypes.wintypes import HANDLE
from ctypes.wintypes import BOOL
from ctypes.wintypes import HWND
from ctypes.wintypes import DWORD
from ctypes.wintypes import WORD
from ctypes.wintypes import LONG
from ctypes.wintypes import ULONG
from ctypes.wintypes import LPCSTR
from ctypes.wintypes import HKEY
from ctypes.wintypes import BYTE

NULL = 0
HDEVINFO = ctypes.c_int
PCTSTR = ctypes.c_char_p
CHAR = ctypes.c_char
LPDWORD = PDWORD = ctypes.POINTER(DWORD)
#~ LPBYTE = PBYTE = ctypes.POINTER(BYTE)
LPBYTE = PBYTE = ctypes.c_void_p        # XXX avoids error about types
PHKEY = ctypes.POINTER(HKEY)

ACCESS_MASK = DWORD
REGSAM = ACCESS_MASK

def byte_buffer(length):
    """Get a buffer for a string"""
    return (BYTE*length)()

def string(buffer):
    s = []
    for c in buffer:
        if c == 0: break
        s.append(chr(c))
    return ''.join(s)

class GUID(ctypes.Structure):
    _fields_ = [
        ('Data1', ctypes.c_ulong),
        ('Data2', ctypes.c_ushort),
        ('Data3', ctypes.c_ushort),
        ('Data4', ctypes.c_ubyte*8),
    ]
    def __str__(self):
        return "{%08x-%04x-%04x-%s-%s}" % (
            self.Data1,
            self.Data2,
            self.Data3,
            ''.join(["%02x" % d for d in self.Data4[:2]]),
            ''.join(["%02x" % d for d in self.Data4[2:]]),
        )

class SP_DEVINFO_DATA(ctypes.Structure):
    _fields_ = [
        ('cbSize', DWORD),
        ('ClassGuid', GUID),
        ('DevInst', DWORD),
        ('Reserved', ULONG_PTR),
    ]
    def __str__(self):
        return "ClassGuid:%s DevInst:%s" % (self.ClassGuid, self.DevInst)
PSP_DEVINFO_DATA = ctypes.POINTER(SP_DEVINFO_DATA)

class SP_DEVICE_INTERFACE_DATA(ctypes.Structure):
    _fields_ = [
        ('cbSize', DWORD),
        ('InterfaceClassGuid', GUID),
        ('Flags', DWORD),
        ('Reserved', ULONG_PTR),
    ]
    def __str__(self):
        return "InterfaceClassGuid:%s Flags:%s" % (self.InterfaceClassGuid, self.Flags)
PSP_DEVICE_INTERFACE_DATA = ctypes.POINTER(SP_DEVICE_INTERFACE_DATA)

PSP_DEVICE_INTERFACE_DETAIL_DATA = ctypes.c_void_p

setupapi = ctypes.windll.LoadLibrary("setupapi")
SetupDiDestroyDeviceInfoList = setupapi.SetupDiDestroyDeviceInfoList
SetupDiDestroyDeviceInfoList.argtypes = [HDEVINFO]
SetupDiDestroyDeviceInfoList.restype = BOOL

SetupDiGetClassDevs = setupapi.SetupDiGetClassDevsA
SetupDiGetClassDevs.argtypes = [ctypes.POINTER(GUID), PCTSTR, HWND, DWORD]
SetupDiGetClassDevs.restype = ValidHandle #HDEVINFO

SetupDiEnumDeviceInterfaces = setupapi.SetupDiEnumDeviceInterfaces
SetupDiEnumDeviceInterfaces.argtypes = [HDEVINFO, PSP_DEVINFO_DATA, ctypes.POINTER(GUID), DWORD, PSP_DEVICE_INTERFACE_DATA]
SetupDiEnumDeviceInterfaces.restype = BOOL

SetupDiGetDeviceInterfaceDetail = setupapi.SetupDiGetDeviceInterfaceDetailA
SetupDiGetDeviceInterfaceDetail.argtypes = [HDEVINFO, PSP_DEVICE_INTERFACE_DATA, PSP_DEVICE_INTERFACE_DETAIL_DATA, DWORD, PDWORD, PSP_DEVINFO_DATA]
SetupDiGetDeviceInterfaceDetail.restype = BOOL

SetupDiGetDeviceRegistryProperty = setupapi.SetupDiGetDeviceRegistryPropertyA
SetupDiGetDeviceRegistryProperty.argtypes = [HDEVINFO, PSP_DEVINFO_DATA, DWORD, PDWORD, PBYTE, DWORD, PDWORD]
SetupDiGetDeviceRegistryProperty.restype = BOOL

SetupDiOpenDevRegKey = setupapi.SetupDiOpenDevRegKey
SetupDiOpenDevRegKey.argtypes = [HDEVINFO, PSP_DEVINFO_DATA, DWORD, DWORD, DWORD, REGSAM]
SetupDiOpenDevRegKey.restype = HKEY

advapi32 = ctypes.windll.LoadLibrary("Advapi32")
RegCloseKey = advapi32.RegCloseKey
RegCloseKey.argtypes = [HKEY]
RegCloseKey.restype = LONG

RegQueryValueEx = advapi32.RegQueryValueExA
RegQueryValueEx.argtypes = [HKEY, LPCSTR, LPDWORD, LPDWORD, LPBYTE, LPDWORD]
RegQueryValueEx.restype = LONG

GUID_CLASS_COMPORT = GUID(0x86e0d1e0L, 0x8089, 0x11d0,
    (ctypes.c_ubyte*8)(0x9c, 0xe4, 0x08, 0x00, 0x3e, 0x30, 0x1f, 0x73))

DIGCF_PRESENT = 2
DIGCF_DEVICEINTERFACE = 16
INVALID_HANDLE_VALUE = 0
ERROR_INSUFFICIENT_BUFFER = 122
SPDRP_HARDWAREID = 1
SPDRP_FRIENDLYNAME = 12
ERROR_NO_MORE_ITEMS = 259
DICS_FLAG_GLOBAL = 1
DIREG_DEV = 0x00000001
KEY_READ = 0x20019
REG_SZ = 1

# workaround for compatibility between Python 2.x and 3.x
PortName = serial.to_bytes([80, 111, 114, 116, 78, 97, 109, 101]) # "PortName"

def comports():
    """This generator scans the device registry for com ports and yields port, desc, hwid"""
    g_hdi = SetupDiGetClassDevs(ctypes.byref(GUID_CLASS_COMPORT), None, NULL, DIGCF_PRESENT|DIGCF_DEVICEINTERFACE);
    #~ for i in range(256):
    for dwIndex in range(256):
        did = SP_DEVICE_INTERFACE_DATA()
        did.cbSize = ctypes.sizeof(did)

        if not SetupDiEnumDeviceInterfaces(g_hdi, None, ctypes.byref(GUID_CLASS_COMPORT), dwIndex, ctypes.byref(did)):
            if ctypes.GetLastError() != ERROR_NO_MORE_ITEMS:
                raise ctypes.WinError()
            break

        dwNeeded = DWORD()
        # get the size
        if not SetupDiGetDeviceInterfaceDetail(g_hdi, ctypes.byref(did), None, 0, ctypes.byref(dwNeeded), None):
            # Ignore ERROR_INSUFFICIENT_BUFFER
            if ctypes.GetLastError() != ERROR_INSUFFICIENT_BUFFER:
                raise ctypes.WinError()
        # allocate buffer
        class SP_DEVICE_INTERFACE_DETAIL_DATA_A(ctypes.Structure):
            _fields_ = [
                ('cbSize', DWORD),
                ('DevicePath', CHAR*(dwNeeded.value - ctypes.sizeof(DWORD))),
            ]
            def __str__(self):
                return "DevicePath:%s" % (self.DevicePath,)
        idd = SP_DEVICE_INTERFACE_DETAIL_DATA_A()
        idd.cbSize = 5
        devinfo = SP_DEVINFO_DATA()
        devinfo.cbSize = ctypes.sizeof(devinfo)
        if not SetupDiGetDeviceInterfaceDetail(g_hdi, ctypes.byref(did), ctypes.byref(idd), dwNeeded, None, ctypes.byref(devinfo)):
            raise ctypes.WinError()

        # hardware ID
        szHardwareID = byte_buffer(250)
        if not SetupDiGetDeviceRegistryProperty(g_hdi, ctypes.byref(devinfo), SPDRP_HARDWAREID, None, ctypes.byref(szHardwareID), ctypes.sizeof(szHardwareID) - 1, None):
            # Ignore ERROR_INSUFFICIENT_BUFFER
            if GetLastError() != ERROR_INSUFFICIENT_BUFFER:
                raise ctypes.WinError()

        # friendly name
        szFriendlyName = byte_buffer(250)
        if not SetupDiGetDeviceRegistryProperty(g_hdi, ctypes.byref(devinfo), SPDRP_FRIENDLYNAME, None, ctypes.byref(szFriendlyName), ctypes.sizeof(szFriendlyName) - 1, None):
            # Ignore ERROR_INSUFFICIENT_BUFFER
            if ctypes.GetLastError() != ERROR_INSUFFICIENT_BUFFER:
                #~ raise IOError("failed to get details for %s (%s)" % (devinfo, szHardwareID.value))
                port_name = None
        else:
            # the real com port name has to read differently...
            hkey = SetupDiOpenDevRegKey(g_hdi, ctypes.byref(devinfo), DICS_FLAG_GLOBAL, 0, DIREG_DEV, KEY_READ)
            port_name_buffer = byte_buffer(250)
            port_name_length = ULONG(ctypes.sizeof(port_name_buffer))
            RegQueryValueEx(hkey, PortName, None, None, ctypes.byref(port_name_buffer), ctypes.byref(port_name_length))
            RegCloseKey(hkey)
            yield string(port_name_buffer), string(szFriendlyName), string(szHardwareID)

    SetupDiDestroyDeviceInfoList(g_hdi)

# test
if __name__ == '__main__':
    import serial

    for port, desc, hwid in sorted(comports()):
        print "%s: %s [%s]" % (port, desc, hwid)
<commit_msg>fix HDEVINFO type<commit_after>import ctypes
import re

def ValidHandle(value):
    if value == 0:
        raise ctypes.WinError()
    return value

import serial
from serial.win32 import ULONG_PTR
from ctypes.wintypes import HANDLE
from ctypes.wintypes import BOOL
from ctypes.wintypes import HWND
from ctypes.wintypes import DWORD
from ctypes.wintypes import WORD
from ctypes.wintypes import LONG
from ctypes.wintypes import ULONG
from ctypes.wintypes import LPCSTR
from ctypes.wintypes import HKEY
from ctypes.wintypes import BYTE

NULL = 0
HDEVINFO = ctypes.c_void_p
PCTSTR = ctypes.c_char_p
CHAR = ctypes.c_char
LPDWORD = PDWORD = ctypes.POINTER(DWORD)
#~ LPBYTE = PBYTE = ctypes.POINTER(BYTE)
LPBYTE = PBYTE = ctypes.c_void_p        # XXX avoids error about types
PHKEY = ctypes.POINTER(HKEY)

ACCESS_MASK = DWORD
REGSAM = ACCESS_MASK

def byte_buffer(length):
    """Get a buffer for a string"""
    return (BYTE*length)()

def string(buffer):
    s = []
    for c in buffer:
        if c == 0: break
        s.append(chr(c))
    return ''.join(s)

class GUID(ctypes.Structure):
    _fields_ = [
        ('Data1', ctypes.c_ulong),
        ('Data2', ctypes.c_ushort),
        ('Data3', ctypes.c_ushort),
        ('Data4', ctypes.c_ubyte*8),
    ]
    def __str__(self):
        return "{%08x-%04x-%04x-%s-%s}" % (
            self.Data1,
            self.Data2,
            self.Data3,
            ''.join(["%02x" % d for d in self.Data4[:2]]),
            ''.join(["%02x" % d for d in self.Data4[2:]]),
        )

class SP_DEVINFO_DATA(ctypes.Structure):
    _fields_ = [
        ('cbSize', DWORD),
        ('ClassGuid', GUID),
        ('DevInst', DWORD),
        ('Reserved', ULONG_PTR),
    ]
    def __str__(self):
        return "ClassGuid:%s DevInst:%s" % (self.ClassGuid, self.DevInst)
PSP_DEVINFO_DATA = ctypes.POINTER(SP_DEVINFO_DATA)

class SP_DEVICE_INTERFACE_DATA(ctypes.Structure):
    _fields_ = [
        ('cbSize', DWORD),
        ('InterfaceClassGuid', GUID),
        ('Flags', DWORD),
        ('Reserved', ULONG_PTR),
    ]
    def __str__(self):
        return "InterfaceClassGuid:%s Flags:%s" % (self.InterfaceClassGuid, self.Flags)
PSP_DEVICE_INTERFACE_DATA = ctypes.POINTER(SP_DEVICE_INTERFACE_DATA)

PSP_DEVICE_INTERFACE_DETAIL_DATA = ctypes.c_void_p

setupapi = ctypes.windll.LoadLibrary("setupapi")
SetupDiDestroyDeviceInfoList = setupapi.SetupDiDestroyDeviceInfoList
SetupDiDestroyDeviceInfoList.argtypes = [HDEVINFO]
SetupDiDestroyDeviceInfoList.restype = BOOL

SetupDiGetClassDevs = setupapi.SetupDiGetClassDevsA
SetupDiGetClassDevs.argtypes = [ctypes.POINTER(GUID), PCTSTR, HWND, DWORD]
SetupDiGetClassDevs.restype = ValidHandle #HDEVINFO

SetupDiEnumDeviceInterfaces = setupapi.SetupDiEnumDeviceInterfaces
SetupDiEnumDeviceInterfaces.argtypes = [HDEVINFO, PSP_DEVINFO_DATA, ctypes.POINTER(GUID), DWORD, PSP_DEVICE_INTERFACE_DATA]
SetupDiEnumDeviceInterfaces.restype = BOOL

SetupDiGetDeviceInterfaceDetail = setupapi.SetupDiGetDeviceInterfaceDetailA
SetupDiGetDeviceInterfaceDetail.argtypes = [HDEVINFO, PSP_DEVICE_INTERFACE_DATA, PSP_DEVICE_INTERFACE_DETAIL_DATA, DWORD, PDWORD, PSP_DEVINFO_DATA]
SetupDiGetDeviceInterfaceDetail.restype = BOOL

SetupDiGetDeviceRegistryProperty = setupapi.SetupDiGetDeviceRegistryPropertyA
SetupDiGetDeviceRegistryProperty.argtypes = [HDEVINFO, PSP_DEVINFO_DATA, DWORD, PDWORD, PBYTE, DWORD, PDWORD]
SetupDiGetDeviceRegistryProperty.restype = BOOL

SetupDiOpenDevRegKey = setupapi.SetupDiOpenDevRegKey
SetupDiOpenDevRegKey.argtypes = [HDEVINFO, PSP_DEVINFO_DATA, DWORD, DWORD, DWORD, REGSAM]
SetupDiOpenDevRegKey.restype = HKEY

advapi32 = ctypes.windll.LoadLibrary("Advapi32")
RegCloseKey = advapi32.RegCloseKey
RegCloseKey.argtypes = [HKEY]
RegCloseKey.restype = LONG

RegQueryValueEx = advapi32.RegQueryValueExA
RegQueryValueEx.argtypes = [HKEY, LPCSTR, LPDWORD, LPDWORD, LPBYTE, LPDWORD]
RegQueryValueEx.restype = LONG

GUID_CLASS_COMPORT = GUID(0x86e0d1e0L, 0x8089, 0x11d0,
    (ctypes.c_ubyte*8)(0x9c, 0xe4, 0x08, 0x00, 0x3e, 0x30, 0x1f, 0x73))

DIGCF_PRESENT = 2
DIGCF_DEVICEINTERFACE = 16
INVALID_HANDLE_VALUE = 0
ERROR_INSUFFICIENT_BUFFER = 122
SPDRP_HARDWAREID = 1
SPDRP_FRIENDLYNAME = 12
ERROR_NO_MORE_ITEMS = 259
DICS_FLAG_GLOBAL = 1
DIREG_DEV = 0x00000001
KEY_READ = 0x20019
REG_SZ = 1

# workaround for compatibility between Python 2.x and 3.x
PortName = serial.to_bytes([80, 111, 114, 116, 78, 97, 109, 101]) # "PortName"

def comports():
    """This generator scans the device registry for com ports and yields port, desc, hwid"""
    g_hdi = SetupDiGetClassDevs(ctypes.byref(GUID_CLASS_COMPORT), None, NULL, DIGCF_PRESENT|DIGCF_DEVICEINTERFACE);
    #~ for i in range(256):
    for dwIndex in range(256):
        did = SP_DEVICE_INTERFACE_DATA()
        did.cbSize = ctypes.sizeof(did)

        if not SetupDiEnumDeviceInterfaces(g_hdi, None, ctypes.byref(GUID_CLASS_COMPORT), dwIndex, ctypes.byref(did)):
            if ctypes.GetLastError() != ERROR_NO_MORE_ITEMS:
                raise ctypes.WinError()
            break

        dwNeeded = DWORD()
        # get the size
        if not SetupDiGetDeviceInterfaceDetail(g_hdi, ctypes.byref(did), None, 0, ctypes.byref(dwNeeded), None):
            # Ignore ERROR_INSUFFICIENT_BUFFER
            if ctypes.GetLastError() != ERROR_INSUFFICIENT_BUFFER:
                raise ctypes.WinError()
        # allocate buffer
        class SP_DEVICE_INTERFACE_DETAIL_DATA_A(ctypes.Structure):
            _fields_ = [
                ('cbSize', DWORD),
                ('DevicePath', CHAR*(dwNeeded.value - ctypes.sizeof(DWORD))),
            ]
            def __str__(self):
                return "DevicePath:%s" % (self.DevicePath,)
        idd = SP_DEVICE_INTERFACE_DETAIL_DATA_A()
        idd.cbSize = 5
        devinfo = SP_DEVINFO_DATA()
        devinfo.cbSize = ctypes.sizeof(devinfo)
        if not SetupDiGetDeviceInterfaceDetail(g_hdi, ctypes.byref(did), ctypes.byref(idd), dwNeeded, None, ctypes.byref(devinfo)):
            raise ctypes.WinError()

        # hardware ID
        szHardwareID = byte_buffer(250)
        if not SetupDiGetDeviceRegistryProperty(g_hdi, ctypes.byref(devinfo), SPDRP_HARDWAREID, None, ctypes.byref(szHardwareID), ctypes.sizeof(szHardwareID) - 1, None):
            # Ignore ERROR_INSUFFICIENT_BUFFER
            if GetLastError() != ERROR_INSUFFICIENT_BUFFER:
                raise ctypes.WinError()

        # friendly name
        szFriendlyName = byte_buffer(250)
        if not SetupDiGetDeviceRegistryProperty(g_hdi, ctypes.byref(devinfo), SPDRP_FRIENDLYNAME, None, ctypes.byref(szFriendlyName), ctypes.sizeof(szFriendlyName) - 1, None):
            # Ignore ERROR_INSUFFICIENT_BUFFER
            if ctypes.GetLastError() != ERROR_INSUFFICIENT_BUFFER:
                #~ raise IOError("failed to get details for %s (%s)" % (devinfo, szHardwareID.value))
                port_name = None
        else:
            # the real com port name has to read differently...
            hkey = SetupDiOpenDevRegKey(g_hdi, ctypes.byref(devinfo), DICS_FLAG_GLOBAL, 0, DIREG_DEV, KEY_READ)
            port_name_buffer = byte_buffer(250)
            port_name_length = ULONG(ctypes.sizeof(port_name_buffer))
            RegQueryValueEx(hkey, PortName, None, None, ctypes.byref(port_name_buffer), ctypes.byref(port_name_length))
            RegCloseKey(hkey)
            yield string(port_name_buffer), string(szFriendlyName), string(szHardwareID)

    SetupDiDestroyDeviceInfoList(g_hdi)

# test
if __name__ == '__main__':
    import serial

    for port, desc, hwid in sorted(comports()):
        print "%s: %s [%s]" % (port, desc, hwid)
<|endoftext|>