<commit_before>#!/usr/bin/env python
#
## @file
#
# A parser for outputting OpenGL. This is based in part on this
# example: https://gist.github.com/deepankarsharma/3494203
#
# This works for most parts, but it is not known to be completely
# conformant with the LDraw specification defined here:
#
# http://www.ldraw.org/article/218.html
#
# In particular it has at least the following issues:
# 1. Optional lines are ignored.
# 2. Culling is disabled as I could not figure out how to get it
#    to work properly.
# 3. Does not support fixed colors.
#
# Hazen 07/14
#

import numpy

from OpenGL import arrays, GL

import datFileParser

# Variables

# Vertex Shader
vertex = """
#version 330
in vec3 vin_position;
uniform mat4 MVP;

void main(void)
{
    vec4 v = vec4(vin_position, 1.0);
    gl_Position = MVP * v;
}
"""

# Fragment Shader
fragment = """
#version 330
uniform vec4 color;
out vec4 fout_color;

void main(void)
{
    fout_color = color;
}
"""

# Functions.

## checkColor
#
# Prints a warning if the color is not an expected value (16 or 24).
#
def checkColor(parsed_line):
    color = int(parsed_line[1])
    if (color != 24) and (color != 16):
        print "Got unexpected color:", color

# Classes.

## GLParser
#
# Parser for creating GL objects. This is used to parse a DAT file and create
# the necessary GL to be able to render the DAT file. This is also a GL object
# so you can draw it in a GL context by calling the render() function. When
# you are finished with it you need to call freeGL() to free the memory associated
# with this object.
#
class GLParser(datFileParser.Parser):
    gl_shader = None

    ## __init__
    #
    # @param matrix (Optional) A 4x4 numpy matrix to use as the transformation matrix for this file. If this is None then the identity matrix is used.
    # @param invert_winding (Optional) Specify if the winding direction should be inverted. Default is False.
    #
    def __init__(self, matrix = None, invert_winding = False):
        datFileParser.Parser.__init__(self, None, None)

        self.cw_winding = False
        self.children = []
        self.depth = 0
        self.invert_next = False
        self.invert_winding = invert_winding
        self.lines_only = False  # This is mostly for debugging.
        self.matrix = matrix
        self.vao_lines = GLVao(GL.GL_LINES)
        self.vao_triangles = GLVao(GL.GL_TRIANGLES)

        if GLParser.gl_shader is None:
            GLParser.gl_shader = GLShader(vertex, fragment)

        if self.matrix is None:
            self.matrix = numpy.identity(4)

    ## addTriangle
    #
    # Add a triangle, winding is always counter-clockwise.
    #
    # @param p1 1st vertex.
    # @param p2 2nd vertex.
    # @param p3 3rd vertex.
    #
    def addTriangle(self, p1, p2, p3):

        if self.lines_only:
            self.vao_lines.addVertex(p1)
            self.vao_lines.addVertex(p2)
            self.vao_lines.addVertex(p2)
            self.vao_lines.addVertex(p3)
            self.vao_lines.addVertex(p3)
            self.vao_lines.addVertex(p1)

        else:
            self.vao_triangles.addVertex(p1)
            self.vao_triangles.addVertex(p2)
            self.vao_triangles.addVertex(p3)

    ## command
    #
    # @param parsed_line A list containing the contents of one line of a parsed file.
    #
    def command(self, parsed_line):
        if (len(parsed_line) > 1):

            # FIXME:
            #  This doesn't work properly in that if you enable face culling then
            #  things will look a bit strange.
            #
            # Handle winding commands.
            if (parsed_line[1] == "BFC"):
                if (parsed_line[2] == "INVERTNEXT"):
                    self.invert_next = True
                elif (parsed_line[2] == "CERTIFY"):
                    if (len(parsed_line) == 3):
                        self.cw_winding = False
                    else:
                        if (parsed_line[3] == "CCW"):
                            self.cw_winding = False
                        else:
                            self.cw_winding = True
                if self.invert_winding:
                    self.cw_winding = not self.cw_winding

                #print "  ", self.depth, parsed_line, self.cw_winding

    ## endFile
    #
    def endFile(self):
        self.vao_lines.finalize(GLParser.gl_shader)
        self.vao_triangles.finalize(GLParser.gl_shader)

    ## freeGL
    #
    # Call this when you are done with this object.
    #
    def freeGL(self):
        if (self.vao_triangles.size > 0):
            self.vao_triangles.freeBuffers()
        if (self.vao_lines.size > 0):
            self.vao_lines.freeBuffers()

        for child in self.children:
            child.freeGL()

    ## line
    #
    # @param parsed_line A list containing the contents of one line of a parsed file.
    #
    def line(self, parsed_line):
        checkColor(parsed_line)
        p1 = self.parsePoint(parsed_line[2:5])
        p2 = self.parsePoint(parsed_line[5:8])
        self.vao_lines.addVertex(p1)
        self.vao_lines.addVertex(p2)

    ## newFile
    #
    # @param parsed_line A list containing the contents of one line of a parsed file.
    #
    # @return A Parser object.
    #
    def newFile(self, parsed_line):
        #print "newFile", self.depth, parsed_line[-1]

        # Parse transformation matrix.
        [x, y, z, a, b, c, d, e, f, g, h, i] = map(float, parsed_line[2:14])
        matrix = numpy.array([[  a,   b,   c,   x],
                              [  d,   e,   f,   y],
                              [  g,   h,   i,   z],
                              [0.0, 0.0, 0.0, 1.0]])
        matrix = numpy.dot(self.matrix, matrix)

        # Figure out windings.
        if self.invert_next:
            invert_winding = not self.invert_winding
        else:
            invert_winding = self.invert_winding

        child = GLParser(matrix = matrix,
                         invert_winding = invert_winding)
        self.invert_next = False
        self.children.append(child)
        return child

    ## optionalLine
    #
    # @param parsed_line A list containing the contents of one line of a parsed file.
    #
    def optionalLine(self, parsed_line):
        checkColor(parsed_line)
        #self.line(parsed_line)

    ## parsePoint
    #
    # Parses point text, applies the current transform matrix and returns the
    # transformed point.
    #
    # @param point A 3 element array containing the points text [x, y, z].
    #
    # @returns A 3 element array of floats.
    #
    def parsePoint(self, point):
        point.append("1.0")
        pi = numpy.array(map(float, point))
        pf = numpy.dot(self.matrix, pi)
        return pf.tolist()[:-1]

    ## quadrilateral
    #
    # @param parsed_line A list containing the contents of one line of a parsed file.
    #
    def quadrilateral(self, parsed_line):
        checkColor(parsed_line)
        p1 = self.parsePoint(parsed_line[2:5])
        p2 = self.parsePoint(parsed_line[5:8])
        p3 = self.parsePoint(parsed_line[8:11])
        p4 = self.parsePoint(parsed_line[11:14])

        if self.lines_only:
            self.vao_lines.addVertex(p1)
            self.vao_lines.addVertex(p2)
            self.vao_lines.addVertex(p2)
            self.vao_lines.addVertex(p3)
            self.vao_lines.addVertex(p3)
            self.vao_lines.addVertex(p4)
            self.vao_lines.addVertex(p4)
            self.vao_lines.addVertex(p1)

        else:
            if self.cw_winding:
                self.addTriangle(p1,p2,p3)
                self.addTriangle(p1,p3,p4)
            else:
                self.addTriangle(p1,p3,p2)
                self.addTriangle(p1,p4,p3)

    ## render
    #
    # Draw the object and all its child objects.
    #
    # @param mvp The model - view - projection matrix or None.
    # @param face_color [r, g, b, a] - Color (0.0 - 1.0) range.
    # @param edge_color [r, g, b, a] - Color (0.0 - 1.0) range.
    #
    def render(self, mvp, face_color, edge_color):

        # Draw object.
        GL.glUseProgram(GLParser.gl_shader.program_id)

        color_id = GLParser.gl_shader.uniform_location('color')
        matrix_id = GLParser.gl_shader.uniform_location('MVP')

        GL.glUniformMatrix4fv(matrix_id, 1, GL.GL_FALSE, mvp)

        if not self.lines_only:
            if (self.vao_triangles.size > 0):
                GL.glUniform4fv(color_id, 1, numpy.array(face_color, dtype=numpy.float32))
                GL.glBindVertexArray(self.vao_triangles.gl_id)
                GL.glDrawArrays(self.vao_triangles.gl_type, 0, self.vao_triangles.size)

        if (self.vao_lines.size > 0):
            GL.glUniform4fv(color_id, 1, numpy.array(edge_color, dtype=numpy.float32))
            GL.glBindVertexArray(self.vao_lines.gl_id)
            GL.glDrawArrays(self.vao_lines.gl_type, 0, self.vao_lines.size)

        GL.glBindVertexArray(0)
        GL.glUseProgram(0)

        # Draw children.
        for child in self.children:
            child.render(mvp, face_color, edge_color)

    ## startFile
    #
    # @param depth The current recursion depth.
    #
    def startFile(self, depth):
        self.depth = depth

    ## triangle
    #
    # @param parsed_line A list containing the contents of one line of a parsed file.
    #
    def triangle(self, parsed_line):
        checkColor(parsed_line)
        p1 = self.parsePoint(parsed_line[2:5])
        p2 = self.parsePoint(parsed_line[5:8])
        p3 = self.parsePoint(parsed_line[8:11])

        if self.cw_winding:
            self.addTriangle(p1,p2,p3)
        else:
            self.addTriangle(p1,p3,p2)

## GLShader
#
# Helper class for using GLSL shader programs.
#
class GLShader(object):

    ## __init__
    #
    # @param vertex A string containing the shader source code for the vertex shader.
    # @param fragment A string containing the shader source code for the fragment shader.
    #
    def __init__(self, vertex, fragment):
        self.program_id = GL.glCreateProgram()
        vs_id = self.add_shader(vertex, GL.GL_VERTEX_SHADER)
        frag_id = self.add_shader(fragment, GL.GL_FRAGMENT_SHADER)

        GL.glAttachShader(self.program_id, vs_id)
        GL.glAttachShader(self.program_id, frag_id)
        GL.glLinkProgram(self.program_id)

        if (GL.glGetProgramiv(self.program_id, GL.GL_LINK_STATUS) != GL.GL_TRUE):
            info = GL.glGetProgramInfoLog(self.program_id)
            GL.glDeleteProgram(self.program_id)
            GL.glDeleteShader(vs_id)
            GL.glDeleteShader(frag_id)
            raise RuntimeError('Error linking program: %s' % (info))
        GL.glDeleteShader(vs_id)
        GL.glDeleteShader(frag_id)

    ## add_shader
    #
    # Helper function for compiling a GLSL shader.
    #
    # @param source String containing shader source code.
    # @param shader_type A OpenGL shader type.
    #
    def add_shader(self, source, shader_type):
        try:
            shader_id = GL.glCreateShader(shader_type)
            GL.glShaderSource(shader_id, source)
            GL.glCompileShader(shader_id)
            if (GL.glGetShaderiv(shader_id, GL.GL_COMPILE_STATUS) != GL.GL_TRUE):
                info = GL.glGetShaderInfoLog(shader_id)
                raise RuntimeError('Shader compilation failed: %s' % (info))
            return shader_id
        except:
            GL.glDeleteShader(shader_id)
            raise

    ## uniform_location
    #
    # Helper function to get location of an OpenGL uniform variable.
    #
    # @param name Name of the variable for which location is to be returned.
    #
    # @return Integer describing location.
    #
    def uniform_location(self, name):
        return GL.glGetUniformLocation(self.program_id, name)

    ## attribute_location
    #
    # Helper function to get location of an OpenGL attribute variable.
    #
    # @param name Name of the variable for which location is to be returned.
    #
    # @return Integer describing location.
    #
    def attribute_location(self, name):
        return GL.glGetAttribLocation(self.program_id, name)

## GLVao
#
# Encapsulates a OpenGL Vertex Array Object.
#
class GLVao(object):

    ## __init__
    #
    # @param gl_type The OpenGL object type.
    # @param color (Optional) Defaults to None.
    #
    def __init__(self, gl_type, color = None):

        self.color = color
        self.gl_type = gl_type

        self.gl_id = 0
        self.size = 0
        self.vertices = []
        self.vbo_id = 0

    ## addVertex.
    #
    # @param vertex [x, y, z] location of the vertex.
    #
    def addVertex(self, vertex):
        self.vertices.extend(vertex)
        self.size += 3

    ## finalize
    #
    # This is called once we have all the vertices to create the necessary OpenGL objects.
    #
    # @param gl_shader A GLShader object.
    #
    def finalize(self, gl_shader):
        if (self.size == 0):
            return

        #print self.vertices, self.n_components

        vertex_data = numpy.array(self.vertices, dtype = numpy.float32)

        self.gl_id = GL.glGenVertexArrays(1)
        GL.glBindVertexArray(self.gl_id)

        self.vbo_id = GL.glGenBuffers(1)
        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.vbo_id)
        GL.glBufferData(GL.GL_ARRAY_BUFFER,
                        arrays.ArrayDatatype.arrayByteCount(vertex_data),
                        vertex_data,
                        GL.GL_STATIC_DRAW)
        GL.glVertexAttribPointer(gl_shader.attribute_location('vin_position'),
                                 3,
                                 GL.GL_FLOAT,
                                 GL.GL_FALSE,
                                 0,
                                 None)
        GL.glEnableVertexAttribArray(0)

        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, 0)
        GL.glBindVertexArray(0)

    ## freeBuffers
    #
    # Frees the OpenGL buffers used by this object.
    #
    def freeBuffers(self):
        GL.glDeleteBuffers(1, [self.vbo_id])
        GL.glDeleteVertexArrays(1, [self.gl_id])
        self.size = 0

    ## getColor
    #
    # @return The color (if any) to use when rendering.
    #
    def getColor(self):
        return self.color

#
# The MIT License
#
# Copyright (c) 2014 Hazen Babcock
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
<commit_msg>Initial work on improving handling of colors.<commit_after>#!/usr/bin/env python
"""
A parser for outputting OpenGL. This is based in part on this
example: https://gist.github.com/deepankarsharma/3494203

This works for most parts, but it is not known to be completely
conformant with the LDraw specification defined here:

http://www.ldraw.org/article/218.html

In particular it has at least the following issues:
1. Optional lines are ignored.
2. Culling is disabled as I could not figure out how to get it
   to work properly.

It is not designed to be fast as it is used mostly for static
rendering.

Hazen 11/15

"""

import numpy

from OpenGL import arrays, GL

import datFileParser

# Variables

# Vertex shader
vertex = """
#version 330
in vec4 vin_position;
in vec4 vin_color;
out vec4 vout_color;
uniform mat4 MVP;

void main(void)
{
    vout_color = vin_color;
    gl_Position = MVP * vin_position;
}
"""

# Fragment shader
fragment = """
#version 330
in vec4 vout_color;
out vec4 fout_color;

void main(void)
{
    fout_color = vout_color;
}
"""

class GLParserException(Exception):
    def __init__(self, message):
        Exception.__init__(self, message)

class GLParser(datFileParser.Parser):
    """
    Parser for creating GL objects. This is used to parse a DAT file and create
    the necessary GL to be able to render the DAT file. This is also a GL object
    so you can draw it in a GL context by calling the render() function. When
    you are finished with it you need to call freeGL() to free the memory associated
    with this object.
    """

    gl_shader = None

    def __init__(self, face_color, edge_color, matrix = None, invert_winding = False):
        datFileParser.Parser.__init__(self, None, None)

        self.cw_winding = False
        self.children = []
        self.depth = 0
        self.edge_color = edge_color
        self.face_color = face_color
        self.invert_next = False
        self.invert_winding = invert_winding
        self.lines_only = False  # This is mostly for debugging.
        self.matrix = matrix
        self.vao_lines = GLVao(GL.GL_LINES)
        self.vao_triangles = GLVao(GL.GL_TRIANGLES)

        if GLParser.gl_shader is None:
            GLParser.gl_shader = GLShader(vertex, fragment)

        if self.matrix is None:
            self.matrix = numpy.identity(4)

    def addLine(self, p1, p2):
        self.vao_lines.addVertex(p1)
        self.vao_lines.addColor(self.edge_color)
        self.vao_lines.addVertex(p2)
        self.vao_lines.addColor(self.edge_color)

    def addTriangle(self, p1, p2, p3):
        if self.lines_only:
            self.addLine(p1, p2)
            self.addLine(p2, p3)
            self.addLine(p3, p1)
        else:
            self.vao_triangles.addVertex(p1)
            self.vao_triangles.addColor(self.face_color)
            self.vao_triangles.addVertex(p2)
            self.vao_triangles.addColor(self.face_color)
            self.vao_triangles.addVertex(p3)
            self.vao_triangles.addColor(self.face_color)

    def command(self, parsed_line):
        if (len(parsed_line) > 1):

            # FIXME:
            #  This doesn't work properly in that if you enable face culling then
            #  things will look a bit strange.
            #
            # Handle winding commands.
            if (parsed_line[1] == "BFC"):
                if (parsed_line[2] == "INVERTNEXT"):
                    self.invert_next = True
                elif (parsed_line[2] == "CERTIFY"):
                    if (len(parsed_line) == 3):
                        self.cw_winding = False
                    else:
                        if (parsed_line[3] == "CCW"):
                            self.cw_winding = False
                        else:
                            self.cw_winding = True
                if self.invert_winding:
                    self.cw_winding = not self.cw_winding

                #print "  ", self.depth, parsed_line, self.cw_winding

    def endFile(self):
        self.vao_lines.finalize(GLParser.gl_shader)
        self.vao_triangles.finalize(GLParser.gl_shader)

    def freeGL(self):
        if (self.vao_triangles.v_size > 0):
            self.vao_triangles.freeBuffers()
        if (self.vao_lines.v_size > 0):
            self.vao_lines.freeBuffers()

        for child in self.children:
            child.freeGL()

    def line(self, parsed_line):
        p1 = self.parsePoint(parsed_line[2:5])
        p2 = self.parsePoint(parsed_line[5:8])
        self.addLine(p1, p2)

    def newFile(self, parsed_line):

        # Parse transformation matrix.
        [x, y, z, a, b, c, d, e, f, g, h, i] = map(float, parsed_line[2:14])
        matrix = numpy.array([[  a,   b,   c,   x],
                              [  d,   e,   f,   y],
                              [  g,   h,   i,   z],
                              [0.0, 0.0, 0.0, 1.0]])
        matrix = numpy.dot(self.matrix, matrix)

        # Figure out windings.
        if self.invert_next:
            invert_winding = not self.invert_winding
        else:
            invert_winding = self.invert_winding

        child = GLParser(self.face_color,
                         self.edge_color,
                         matrix = matrix,
                         invert_winding = invert_winding)
        self.invert_next = False
        self.children.append(child)
        return child

    def optionalLine(self, parsed_line):
        pass

    def parsePoint(self, point):
        point.append("1.0")
        pi = numpy.array(map(float, point))
        pf = numpy.dot(self.matrix, pi)
        return pf.tolist()[:-1] + [1.0]

    def quadrilateral(self, parsed_line):
        p1 = self.parsePoint(parsed_line[2:5])
        p2 = self.parsePoint(parsed_line[5:8])
        p3 = self.parsePoint(parsed_line[8:11])
        p4 = self.parsePoint(parsed_line[11:14])

        if self.lines_only:
            self.addLine(p1, p2)
            self.addLine(p2, p3)
            self.addLine(p3, p4)
            self.addLine(p4, p1)

        else:
            if self.cw_winding:
                self.addTriangle(p1,p2,p3)
                self.addTriangle(p1,p3,p4)
            else:
                self.addTriangle(p1,p3,p2)
                self.addTriangle(p1,p4,p3)

    def render(self, mvp):

        # Draw object.
        GL.glUseProgram(GLParser.gl_shader.program_id)

        matrix_id = GLParser.gl_shader.uniformLocation('MVP')

        GL.glUniformMatrix4fv(matrix_id, 1, GL.GL_FALSE, mvp)

        if not self.lines_only:
            if (self.vao_triangles.v_size > 0):
                GL.glBindVertexArray(self.vao_triangles.gl_id)
                GL.glDrawArrays(self.vao_triangles.gl_type, 0, self.vao_triangles.v_size)

        if (self.vao_lines.v_size > 0):
            GL.glBindVertexArray(self.vao_lines.gl_id)
            GL.glDrawArrays(self.vao_lines.gl_type, 0, self.vao_lines.v_size)

        GL.glBindVertexArray(0)
        GL.glUseProgram(0)

        # Draw children.
        for child in self.children:
            child.render(mvp)

    def startFile(self, depth):
        self.depth = depth

    def triangle(self, parsed_line):
        p1 = self.parsePoint(parsed_line[2:5])
        p2 = self.parsePoint(parsed_line[5:8])
        p3 = self.parsePoint(parsed_line[8:11])

        if self.cw_winding:
            self.addTriangle(p1,p2,p3)
        else:
            self.addTriangle(p1,p3,p2)

class GLShader(object):

    def __init__(self, vertex, fragment):
        self.program_id = GL.glCreateProgram()
        vs_id = self.add_shader(vertex, GL.GL_VERTEX_SHADER)
        frag_id = self.add_shader(fragment, GL.GL_FRAGMENT_SHADER)

        GL.glAttachShader(self.program_id, vs_id)
        GL.glAttachShader(self.program_id, frag_id)
        GL.glLinkProgram(self.program_id)

        if (GL.glGetProgramiv(self.program_id, GL.GL_LINK_STATUS) != GL.GL_TRUE):
            info = GL.glGetProgramInfoLog(self.program_id)
            GL.glDeleteProgram(self.program_id)
            GL.glDeleteShader(vs_id)
            GL.glDeleteShader(frag_id)
            raise RuntimeError('Error linking program: %s' % (info))
        GL.glDeleteShader(vs_id)
        GL.glDeleteShader(frag_id)

    def add_shader(self, source, shader_type):
        try:
            shader_id = GL.glCreateShader(shader_type)
            GL.glShaderSource(shader_id, source)
            GL.glCompileShader(shader_id)
            if (GL.glGetShaderiv(shader_id, GL.GL_COMPILE_STATUS) != GL.GL_TRUE):
                info = GL.glGetShaderInfoLog(shader_id)
                raise RuntimeError('Shader compilation failed: %s' % (info))
            return shader_id
        except:
            GL.glDeleteShader(shader_id)
            raise

    def uniformLocation(self, name):
        return GL.glGetUniformLocation(self.program_id, name)

    def attributeLocation(self, name):
        return GL.glGetAttribLocation(self.program_id, name)

class GLVao(object):

    def __init__(self, gl_type, color = None):

        self.color = color
        self.gl_type = gl_type

        self.c_size = 0
        self.gl_id = 0
        self.vertices = []
        self.colors = []
        self.vbo_id = 0
        self.v_size = 0

    def addColor(self, color):
        self.colors.extend(color)
        self.c_size += 4

    def addVertex(self, vertex):
        self.vertices.extend(vertex)
        self.v_size += 4

    def finalize(self, gl_shader):
        if (self.v_size == 0):
            return

        if (self.c_size != self.v_size):
            raise GLParserException("Number of colors vertices does not match number of position vertices")

        color_data = numpy.array(self.colors, dtype = numpy.float32)
        vertex_data = numpy.array(self.vertices, dtype = numpy.float32)

        self.gl_id = GL.glGenVertexArrays(1)
        GL.glBindVertexArray(self.gl_id)

        # Create buffers.
        self.vbo_id = GL.glGenBuffers(2)

        # Fill first buffer with the vertices.
        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.vbo_id[0])
        GL.glBufferData(GL.GL_ARRAY_BUFFER,
                        arrays.ArrayDatatype.arrayByteCount(vertex_data),
                        vertex_data,
                        GL.GL_STATIC_DRAW)
        GL.glVertexAttribPointer(gl_shader.attributeLocation('vin_position'),
                                 4,
                                 GL.GL_FLOAT,
                                 GL.GL_FALSE,
                                 0,
                                 None)
        GL.glEnableVertexAttribArray(0)

        # Fill second buffer with the colors.
        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.vbo_id[1])
        GL.glBufferData(GL.GL_ARRAY_BUFFER,
                        arrays.ArrayDatatype.arrayByteCount(color_data),
                        color_data,
                        GL.GL_STATIC_DRAW)
        GL.glVertexAttribPointer(gl_shader.attributeLocation('vin_color'),
                                 4,
                                 GL.GL_FLOAT,
                                 GL.GL_FALSE,
                                 0,
                                 None)
        GL.glEnableVertexAttribArray(1)

        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, 0)
        GL.glBindVertexArray(0)

    def freeBuffers(self):
        GL.glDeleteBuffers(2, self.vbo_id)
        GL.glDeleteVertexArrays(1, [self.gl_id])
        self.v_size = 0

    def getColor(self):
        return self.color

#
# The MIT License
#
# Copyright (c) 2015 Hazen Babcock
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
<|endoftext|>


<commit_before>import string

from tests import SimpleTestSchema
from mongo_dynamic_fixture import N
from mongo_dynamic_fixture import G
from mongo_dynamic_fixture.test import MongoTestCase

class FacadesTestCase(MongoTestCase):

    def setUp(self):
        super(FacadesTestCase, self).setUp()
        db_name = 'db_test'
        coll_name = 'coll_test'
        self.conn = self.mongo_client[db_name][coll_name]

    def test_N_no_schema_no_extra(self):
        actual_data = N(_id='123', key_1='value-1', key_2='value-2')
        expected_data = {
            '_id': '123',
            'key_1': 'value-1',
            'key_2': 'value-2'
        }
        self.assertEqual(actual_data, expected_data)

    def test_N_no_schema(self):
        actual_data = N(_id='123', key_1='value-1', key_2='value-2',
                        extra={'key-3': 'value-3'})
        expected_data = {
            '_id': '123',
            'key_1': 'value-1',
            'key_2': 'value-2',
            'key-3': 'value-3'
        }
        self.assertEqual(actual_data, expected_data)

    def test_N_with_schema_no_extra(self):
        actual_data = N(SimpleTestSchema(), _id='123', nest_3__double=999.999)

        self.assertEqual(
            set(actual_data.keys()),
            set(SimpleTestSchema.schema.keys()).union(set(['_id'])))
        self.assertEqual(actual_data['_id'], '123')
        self.assertTrue(0 <= len(actual_data['array']) <= 10)
        self.assertTrue(all([isinstance(i, int)
                             for i in actual_data['array']]))

        self.assertEqual(set(actual_data['nest-1'].keys()),
                         set(SimpleTestSchema.schema['nest-1'].keys()))
        self.assertTrue(0 <= actual_data['nest-1']['integer'] <= 100)

        self.assertTrue(
            0 <= len(actual_data['nest-1']['nest-2']['string']) <= 10)
        self.assertTrue(
            all([s in (string.ascii_letters + string.digits)
                 for s in actual_data['nest-1']['nest-2']['string']]))
        self.assertEqual(actual_data['nest_3']['double'], 999.999)

    def test_N_with_schema(self):
        actual_data = N(SimpleTestSchema(), _id='123', nest_3__double=999.999,
                        extra={'nest-1__nest-2__string': 'abcdef',
                               'nest-1__integer': 10000})

        self.assertEqual(
            set(actual_data.keys()),
            set(SimpleTestSchema.schema.keys()).union(set(['_id'])))
        self.assertEqual(actual_data['_id'], '123')
        self.assertTrue(0 <= len(actual_data['array']) <= 10)
        self.assertTrue(all([isinstance(i, int)
                             for i in actual_data['array']]))

        self.assertEqual(set(actual_data['nest-1'].keys()),
                         set(SimpleTestSchema.schema['nest-1'].keys()))
        self.assertEqual(actual_data['nest-1']['integer'], 10000)

        self.assertEqual(
            set(actual_data['nest-1']['nest-2'].keys()),
            set(SimpleTestSchema.schema['nest-1']['nest-2'].keys()))
        self.assertEqual(actual_data['nest-1']['nest-2']['string'], 'abcdef')
        self.assertEqual(actual_data['nest_3']['double'], 999.999)

    def test_G_no_schema_no_extra(self):
        actual_data = G(self.conn, _id='123', key_1='value-1', key_2='value-2')
        expected_data = {
            '_id': '123',
            'key_1': 'value-1',
            'key_2': 'value-2'
        }
        self.assertEqual(actual_data, expected_data)

        documents = list(self.conn.find())
        self.assertEqual(len(documents), 1)
        self.assertEqual(documents[0], expected_data)

    def test_G_no_schema(self):
        actual_data = G(self.conn, _id='123', key_1='value-1', key_2='value-2',
                        extra={'key-3': 'value-3'})
        expected_data = {
            '_id': '123',
            'key_1': 'value-1',
            'key_2': 'value-2',
            'key-3': 'value-3'
        }
        self.assertEqual(actual_data, expected_data)

        documents = list(self.conn.find())
        self.assertEqual(len(documents), 1)
        self.assertEqual(documents[0], expected_data)

    def test_G_with_schema_no_extra(self):
        actual_data = G(self.conn, SimpleTestSchema(), _id='123',
                        nest_3__double=999.999)

        documents = list(self.conn.find())
        self.assertEqual(len(documents), 1)

        for data in [actual_data, documents[0]]:
            self.assertEqual(
                set(data.keys()),
                set(SimpleTestSchema.schema.keys()).union(set(['_id'])))
            self.assertEqual(data['_id'], '123')
            self.assertTrue(0 <= len(data['array']) <= 10)
            self.assertTrue(all([isinstance(i, int) for i in data['array']]))

            self.assertEqual(set(data['nest-1'].keys()),
                             set(SimpleTestSchema.schema['nest-1'].keys()))
            self.assertTrue(0 <= actual_data['nest-1']['integer'] <= 100)

            self.assertTrue(0 <= len(data['nest-1']['nest-2']['string']) <= 10)
            self.assertTrue(
                all([s in (string.ascii_letters + string.digits)
                     for s in data['nest-1']['nest-2']['string']]))
            self.assertEqual(data['nest_3']['double'], 999.999)

    def test_G_with_schema(self):
        actual_data = G(self.conn, SimpleTestSchema(), _id='123',
                        nest_3__double=999.999,
                        extra={'nest-1__nest-2__string': 'abcdef',
                               'nest-1__integer': 10000})

        documents = list(self.conn.find())
        self.assertEqual(len(documents), 1)

        for data in [actual_data, documents[0]]:
            self.assertEqual(
                set(data.keys()),
                set(SimpleTestSchema.schema.keys()).union(set(['_id'])))
            self.assertEqual(data['_id'], '123')
            self.assertTrue(0 <= len(data['array']) <= 10)
            self.assertTrue(all([isinstance(i, int) for i in data['array']]))

            self.assertEqual(set(data['nest-1'].keys()),
                             set(SimpleTestSchema.schema['nest-1'].keys()))
            self.assertEqual(data['nest-1']['integer'], 10000)

            self.assertEqual(
                set(data['nest-1']['nest-2'].keys()),
                set(SimpleTestSchema.schema['nest-1']['nest-2'].keys()))
            self.assertEqual(data['nest-1']['nest-2']['string'], 'abcdef')
            self.assertEqual(data['nest_3']['double'], 999.999)
<commit_msg>Fixed FacadesTestCase<commit_after>import string

from tests import SimpleTestSchema
from mongo_dynamic_fixture import N
from mongo_dynamic_fixture import G
from mongo_dynamic_fixture.test import MongoTestCase

class FacadesTestCase(MongoTestCase):

    def setUp(self):
        super(FacadesTestCase, self).setUp()
        db_name = 'db_test'
        coll_name = 'coll_test'
        self.conn = self.mongo_client[db_name][coll_name]

    def test_N_no_schema_no_extra(self):
        actual_data = N(_id='123', key_1='value-1', key_2='value-2')
        expected_data = {
            '_id': '123',
            'key_1': 'value-1',
            'key_2': 'value-2'
        }
        self.assertEqual(actual_data, expected_data)

    def test_N_no_schema(self):
        actual_data = N(_id='123', key_1='value-1', key_2='value-2',
                        extra={'key-3': 'value-3'})
        expected_data = {
            '_id': '123',
            'key_1': 'value-1',
            'key_2': 'value-2',
            'key-3': 'value-3'
        }
        self.assertEqual(actual_data, expected_data)

    def test_N_with_schema_no_extra(self):
        actual_data = N(SimpleTestSchema, _id='123', nest_3__double=999.999)

        self.assertEqual(
            set(actual_data.keys()),
            set(SimpleTestSchema.schema.keys()).union(set(['_id'])))
        self.assertEqual(actual_data['_id'], '123')
        self.assertTrue(0 <= len(actual_data['array']) <= 10)
        self.assertTrue(all([isinstance(i, int)
                             for i in actual_data['array']]))

        self.assertEqual(set(actual_data['nest-1'].keys()),
                         set(SimpleTestSchema.schema['nest-1'].keys()))
        self.assertTrue(0 <= actual_data['nest-1']['integer'] <= 100)

        self.assertTrue(
            0 <= len(actual_data['nest-1']['nest-2']['string']) <= 10)
        self.assertTrue(
            all([s in (string.ascii_letters + string.digits)
                 for s in actual_data['nest-1']['nest-2']['string']]))
        self.assertEqual(actual_data['nest_3']['double'], 999.999)

    def test_N_with_schema(self):
        actual_data = N(SimpleTestSchema, _id='123', nest_3__double=999.999,
                        extra={'nest-1__nest-2__string': 'abcdef',
                               'nest-1__integer': 10000})

        self.assertEqual(
            set(actual_data.keys()),
            set(SimpleTestSchema.schema.keys()).union(set(['_id'])))
        self.assertEqual(actual_data['_id'], '123')
        self.assertTrue(0 <= len(actual_data['array']) <= 10)
        self.assertTrue(all([isinstance(i, int)
                             for i in actual_data['array']]))

        self.assertEqual(set(actual_data['nest-1'].keys()),
                         set(SimpleTestSchema.schema['nest-1'].keys()))
        self.assertEqual(actual_data['nest-1']['integer'], 10000)

        self.assertEqual(
            set(actual_data['nest-1']['nest-2'].keys()),
            set(SimpleTestSchema.schema['nest-1']['nest-2'].keys()))
        self.assertEqual(actual_data['nest-1']['nest-2']['string'], 'abcdef')
        self.assertEqual(actual_data['nest_3']['double'], 999.999)

    def test_G_no_schema_no_extra(self):
        actual_data = G(self.conn, _id='123', key_1='value-1', key_2='value-2')
        expected_data = {
            '_id': '123',
            'key_1': 'value-1',
            'key_2': 'value-2'
        }
        self.assertEqual(actual_data, expected_data)

        documents = list(self.conn.find())
        self.assertEqual(len(documents), 1)
        self.assertEqual(documents[0], expected_data)

    def test_G_no_schema(self):
        actual_data = G(self.conn, _id='123', key_1='value-1', key_2='value-2',
                        extra={'key-3': 'value-3'})
        expected_data = {
            '_id': '123',
            'key_1': 'value-1',
            'key_2': 'value-2',
            'key-3': 'value-3'
        }
        self.assertEqual(actual_data, expected_data)

        documents = list(self.conn.find())
        self.assertEqual(len(documents), 1)
        self.assertEqual(documents[0], expected_data)

    def test_G_with_schema_no_extra(self):
        actual_data = G(self.conn, SimpleTestSchema, _id='123',
                        nest_3__double=999.999)

        documents = list(self.conn.find())
        self.assertEqual(len(documents), 1)

        for data in [actual_data, documents[0]]:
            self.assertEqual(
                set(data.keys()),
                set(SimpleTestSchema.schema.keys()).union(set(['_id'])))
            self.assertEqual(data['_id'], '123')
            self.assertTrue(0 <= len(data['array']) <= 10)
            self.assertTrue(all([isinstance(i, int) for i in data['array']]))

            self.assertEqual(set(data['nest-1'].keys()),
                             set(SimpleTestSchema.schema['nest-1'].keys()))
            self.assertTrue(0 <= actual_data['nest-1']['integer'] <= 100)

            self.assertTrue(0 <= len(data['nest-1']['nest-2']['string']) <= 10)
            self.assertTrue(
                all([s in (string.ascii_letters + string.digits)
                     for s in data['nest-1']['nest-2']['string']]))
            self.assertEqual(data['nest_3']['double'], 999.999)

    def test_G_with_schema(self):
        actual_data = G(self.conn, SimpleTestSchema, _id='123',
                        nest_3__double=999.999,
                        extra={'nest-1__nest-2__string': 'abcdef',
                               'nest-1__integer': 10000})

        documents = list(self.conn.find())
        self.assertEqual(len(documents), 1)

        for data in [actual_data, documents[0]]:
            self.assertEqual(
                set(data.keys()),
                set(SimpleTestSchema.schema.keys()).union(set(['_id'])))
            self.assertEqual(data['_id'], '123')
            self.assertTrue(0 <= len(data['array']) <= 10)
            self.assertTrue(all([isinstance(i, int) for i in data['array']]))

            self.assertEqual(set(data['nest-1'].keys()),
                             set(SimpleTestSchema.schema['nest-1'].keys()))
            self.assertEqual(data['nest-1']['integer'], 10000)

            self.assertEqual(
                set(data['nest-1']['nest-2'].keys()),
                set(SimpleTestSchema.schema['nest-1']['nest-2'].keys()))
            self.assertEqual(data['nest-1']['nest-2']['string'], 'abcdef')
            self.assertEqual(data['nest_3']['double'], 999.999)
<|endoftext|>


<commit_before>from django.http import HttpResponseForbidden, HttpResponseNotFound
from django.views.static import serve

from ..core.config import config
from ..utils.auth import has_perm
from ..utils.rest_api import ModelViewSet, ValidationError
from .access_permissions import MediafileAccessPermissions
from .models import Mediafile

# Viewsets for the REST API

class MediafileViewSet(ModelViewSet):
    """
    API endpoint for mediafile objects.

    There are the following views: metadata, list, retrieve, create,
    partial_update, update and destroy.
    """

    access_permissions = MediafileAccessPermissions()
    queryset = Mediafile.objects.all()

    def check_view_permissions(self):
        """
        Returns True if the user has required permissions.
        """
        if self.action in ("list", "retrieve"):
            result = self.get_access_permissions().check_permissions(self.request.user)
        elif self.action == "metadata":
            result = has_perm(self.request.user, "mediafiles.can_see")
        elif self.action == "create":
            result = has_perm(self.request.user, "mediafiles.can_see") and has_perm(
                self.request.user, "mediafiles.can_upload"
            )
        elif self.action in ("partial_update", "update"):
            result = (
                has_perm(self.request.user, "mediafiles.can_see")
                and has_perm(self.request.user, "mediafiles.can_upload")
                and has_perm(self.request.user, "mediafiles.can_manage")
            )
        elif self.action == "destroy":
            result = has_perm(self.request.user, "mediafiles.can_see") and has_perm(
                self.request.user, "mediafiles.can_manage"
            )
        else:
            result = False
        return result

    def create(self, request, *args, **kwargs):
        """
        Customized view endpoint to upload a new file.
        """
        # Check permission to check if the uploader has to be changed.
        uploader_id = self.request.data.get("uploader_id")
        if (
            uploader_id
            and not has_perm(request.user, "mediafiles.can_manage")
            and str(self.request.user.pk) != str(uploader_id)
        ):
            self.permission_denied(request)
        if not self.request.data.get("mediafile"):
            raise ValidationError({"detail": "You forgot to provide a file."})
        return super().create(request, *args, **kwargs)

    def destroy(self, request, *args, **kwargs):
        """
        Customized view endpoint to delete uploaded files.

        Does also delete the file from filesystem.
        """
        # To avoid Django calling save() and triggering autoupdate we do not
        # use the builtin method mediafile.mediafile.delete() but call
        # mediafile.mediafile.storage.delete(...) directly. This may have
        # unattended side effects so be careful especially when accessing files
        # on server via Django methods (file, open(), save(), ...).
        mediafile = self.get_object()
        mediafile.mediafile.storage.delete(mediafile.mediafile.name)
        for logo in config["logos_available"]:
            if config[logo]["path"] == mediafile.mediafile.url:
                config[logo] = {
                    "display_name": config[logo]["display_name"],
                    "path": "",
                }
        return super().destroy(request, *args, **kwargs)

def protected_serve(request, path, document_root=None, show_indexes=False):
    try:
        mediafile = Mediafile.objects.get(mediafile=path)
    except Mediafile.DoesNotExist:
        return HttpResponseNotFound(content="Not found.")

    can_see = has_perm(request.user, "mediafiles.can_see")
    is_special_file = mediafile.is_logo() or mediafile.is_font()
    is_hidden_but_no_perms = mediafile.hidden and not has_perm(
        request.user, "mediafiles.can_see_hidden"
    )

    if not is_special_file and (not can_see or is_hidden_but_no_perms):
        return HttpResponseForbidden(content="Forbidden.")
    else:
        return serve(request, path, document_root, show_indexes)
<commit_msg>Unset font if the mediafile is deleted<commit_after>from django.http import HttpResponseForbidden, HttpResponseNotFound
from django.views.static import serve

from ..core.config import config
from ..utils.auth import has_perm
from ..utils.rest_api import ModelViewSet, ValidationError
from .access_permissions import MediafileAccessPermissions
from .models import Mediafile

# Viewsets for the REST API

class MediafileViewSet(ModelViewSet):
    """
    API endpoint for mediafile objects.

    There are the following views: metadata, list, retrieve, create,
    partial_update, update and destroy.
    """

    access_permissions = MediafileAccessPermissions()
    queryset = Mediafile.objects.all()

    def check_view_permissions(self):
        """
        Returns True if the user has required permissions.
        """
        if self.action in ("list", "retrieve"):
            result = self.get_access_permissions().check_permissions(self.request.user)
        elif self.action == "metadata":
            result = has_perm(self.request.user, "mediafiles.can_see")
        elif self.action == "create":
            result = has_perm(self.request.user, "mediafiles.can_see") and has_perm(
                self.request.user, "mediafiles.can_upload"
            )
        elif self.action in ("partial_update", "update"):
            result = (
                has_perm(self.request.user, "mediafiles.can_see")
                and has_perm(self.request.user, "mediafiles.can_upload")
                and has_perm(self.request.user, "mediafiles.can_manage")
            )
        elif self.action == "destroy":
            result = has_perm(self.request.user, "mediafiles.can_see") and has_perm(
                self.request.user, "mediafiles.can_manage"
            )
        else:
            result = False
        return result

    def create(self, request, *args, **kwargs):
        """
        Customized view endpoint to upload a new file.
        """
        # Check permission to check if the uploader has to be changed.
        uploader_id = self.request.data.get("uploader_id")
        if (
            uploader_id
            and not has_perm(request.user, "mediafiles.can_manage")
            and str(self.request.user.pk) != str(uploader_id)
        ):
            self.permission_denied(request)
        if not self.request.data.get("mediafile"):
            raise ValidationError({"detail": "You forgot to provide a file."})
        return super().create(request, *args, **kwargs)

    def destroy(self, request, *args, **kwargs):
        """
        Customized view endpoint to delete uploaded files.

        Does also delete the file from filesystem.
        """
        # To avoid Django calling save() and triggering autoupdate we do not
        # use the builtin method mediafile.mediafile.delete() but call
        # mediafile.mediafile.storage.delete(...) directly. This may have
        # unattended side effects so be careful especially when accessing files
        # on server via Django methods (file, open(), save(), ...).
        mediafile = self.get_object()
        mediafile.mediafile.storage.delete(mediafile.mediafile.name)

        # check if the file was used as a logo or font
        for logo in config["logos_available"]:
            if config[logo]["path"] == mediafile.mediafile.url:
                config[logo] = {
                    "display_name": config[logo]["display_name"],
                    "path": "",
                }
        for font in config["fonts_available"]:
            if config[font]["path"] == mediafile.mediafile.url:
                config[font] = {
                    "display_name": config[font]["display_name"],
                    "default": config[font]["default"],
                    "path": "",
                }
        return super().destroy(request, *args, **kwargs)

def protected_serve(request, path, document_root=None, show_indexes=False):
    try:
        mediafile = Mediafile.objects.get(mediafile=path)
    except Mediafile.DoesNotExist:
        return HttpResponseNotFound(content="Not found.")

    can_see = has_perm(request.user, "mediafiles.can_see")
    is_special_file = mediafile.is_logo() or mediafile.is_font()
    is_hidden_but_no_perms = mediafile.hidden and not has_perm(
        request.user, "mediafiles.can_see_hidden"
    )

    if not is_special_file and (not can_see or is_hidden_but_no_perms):
        return HttpResponseForbidden(content="Forbidden.")
    else:
        return serve(request, path, document_root, show_indexes)
<|endoftext|>


<commit_before>import numpy as np
from csxtools.fastccd import correct_images, photon_count
from numpy.testing import (assert_array_max_ulp, assert_array_equal,
                           assert_array_almost_equal)

def test_correct_images():
    x = np.ones((3, 10, 10), dtype=np.uint16)
    x[0] = x[0] * 0x0010
    x[1] = x[1] * 0x8020
    x[2] = x[2] * 0xC030

    y = np.ones((3, 10, 10), dtype=np.float32)
    y[0] = y[0] * 0x0010
    y[1] = y[1] * 0x0020
    y[2] = y[2] * 0x0030
    z = correct_images(x, y)
    assert_array_max_ulp(z, np.zeros_like(x))

def test_photon_count():
    x = np.array([[0,  0,  0,  0,  0,  0,  0,  0],
                  [0,  0,  0,  0,  0,  4,  3,  0],
                  [0,  0,  0, 10,  0,  4,  0,  0],
                  [0,  0,  4,  6,  2,  0,  0,  0],
                  [0,  0,  0,  0,  0,  0,  0,  0],
                  [0,  0,  0,  0,  0,  0,  0,  0]], dtype=np.float32)

    y = np.zeros_like(x)
    y[2, 3] = 20
    z = np.zeros_like(x)
    z[2, 3] = np.std(np.array([10, 6, 4, 2, 0, 0, 0, 0, 0],
                              dtype=np.float32)[:3])

    op = photon_count(np.array([x, x, x], dtype=np.float32),
                      thresh=(5, 13), mean_filter=(10, 30), nsum=3)

    assert_array_equal(op[0], np.array([y, y, y]))
    assert_array_almost_equal(op[1], np.array([z, z, z]), decimal=6)
<commit_msg>[ENH] Fixed to use nsum as a parameter<commit_after>import numpy as np
from csxtools.fastccd import correct_images, photon_count
from numpy.testing import (assert_array_max_ulp, assert_array_equal,
                           assert_array_almost_equal)

def test_correct_images():
    x = np.ones((3, 10, 10), dtype=np.uint16)
    x[0] = x[0] * 0x0010
    x[1] = x[1] * 0x8020
    x[2] = x[2] * 0xC030

    y = np.ones((3, 10, 10), dtype=np.float32)
    y[0] = y[0] * 0x0010
    y[1] = y[1] * 0x0020
    y[2] = y[2] * 0x0030
    z = correct_images(x, y)
    assert_array_max_ulp(z, np.zeros_like(x))

def test_photon_count():
    x = np.array([[0,  0,  0,  0,  0,  0,  0,  0],
                  [0,  0,  0,  0,  0,  4,  3,  0],
                  [0,  0,  0, 10,  0,  4,  0,  0],
                  [0,  0,  4,  6,  2,  0,  0,  0],
                  [0,  0,  0,  0,  0,  0,  0,  0],
                  [0,  0,  0,  0,  0,  0,  0,  0]], dtype=np.float32)

    nsum = 3

    y = np.zeros_like(x)
    y[2, 3] = 20
    z = np.zeros_like(x)
    z[2, 3] = np.std(np.array([10, 6, 4, 2, 0, 0, 0, 0, 0],
                              dtype=np.float32)[:nsum])

    op = photon_count(np.array([x, x, x], dtype=np.float32),
                      thresh=(5, 13), mean_filter=(10, 30), nsum=nsum)

    assert_array_equal(op[0], np.array([y, y, y]))
    assert_array_almost_equal(op[1], np.array([z, z, z]), decimal=6)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import json
from unittest import TestCase

import responses
from mock import patch

from frigg_worker.fetcher import fetch_task, start_build

def mock_config(key):
    if key == 'DISPATCHER_URL':
        return 'http://example.com'
    if key == 'DISPATCHER_TOKEN':
        return 'token'

class FetcherTestCase(TestCase):

    @patch('frigg_worker.jobs.Build.run_tests')
    def test_start_build(self, mock_runtests):
        start_build({
            'id': 1,
        }, {})
        self.assertTrue(mock_runtests.called)

    @patch('frigg_worker.fetcher.config', mock_config)
    @responses.activate
    def test_fetch_task(self):

        responses.add(
            responses.GET,
            mock_config('DISPATCHER_URL'),
            body=json.dumps({'job': {'id': 1}}),
            content_type='application/json'
        )

        task = fetch_task('http://example.com', None)
        self.assertEqual(task['id'], 1)

    @patch('frigg_worker.fetcher.config', mock_config)
    @responses.activate
    def test_fetch_task_custom_dispatcher(self):

        responses.add(
            responses.GET,
            'http://dispatcher.io',
            body=json.dumps({'job': {'id': 1}}),
            content_type='application/json'
        )

        task = fetch_task('http://dispatcher.io', dispatcher_token=None)
        self.assertEqual(task['id'], 1)
<commit_msg>Fixes sorting of imports api.py<commit_after># -*- coding: utf-8 -*-
import json
from unittest import TestCase

import responses
import mock
from mock import patch

from frigg_worker.fetcher import fetch_task, start_build

def mock_config(key):
    if key == 'DISPATCHER_URL':
        return 'http://example.com'
    if key == 'DISPATCHER_TOKEN':
        return 'token'

class FetcherTestCase(TestCase):

    @mock.patch('docker.manager.Docker')
    @patch('frigg_worker.jobs.Build.run_tests')
    def test_start_build(self, mock_runtests):
        start_build({
            'id': 1,
        }, {})
        self.assertTrue(mock_runtests.called)

    @patch('frigg_worker.fetcher.config', mock_config)
    @responses.activate
    def test_fetch_task(self):

        responses.add(
            responses.GET,
            mock_config('DISPATCHER_URL'),
            body=json.dumps({'job': {'id': 1}}),
            content_type='application/json'
        )

        task = fetch_task('http://example.com', None)
        self.assertEqual(task['id'], 1)

    @patch('frigg_worker.fetcher.config', mock_config)
    @responses.activate
    def test_fetch_task_custom_dispatcher(self):

        responses.add(
            responses.GET,
            'http://dispatcher.io',
            body=json.dumps({'job': {'id': 1}}),
            content_type='application/json'
        )

        task = fetch_task('http://dispatcher.io', dispatcher_token=None)
        self.assertEqual(task['id'], 1)
<|endoftext|>


<commit_before>import pytest

import vtki
from vtki import examples

datasets = [
    examples.load_uniform(), # UniformGrid
    examples.load_rectilinear(), # RectilinearGrid
    examples.load_hexbean(), # UnstructuredGrid
    examples.load_airplane(), # PolyData
    examples.load_structured(), # StructuredGrid
]
normals = ['x', 'y', '-z', (1,1,1), (3.3, 5.4, 0.8)]

def test_clip_filter():
    """This tests the clip filter on all datatypes avaialble filters"""
    for i, dataset in enumerate(datasets):
        clp = dataset.clip(normal=normals[i], invert=True)
        assert clp is not None
        assert isinstance(clp, vtki.UnstructuredGrid)

def test_slice_filter():
    """This tests the slice filter on all datatypes avaialble filters"""
    for i, dataset in enumerate(datasets):
        slc = dataset.slice(normal=normals[i])
        assert slc is not None
        assert isinstance(slc, vtki.PolyData)

def test_slice_orthogonal_filter():
    """This tests the slice filter on all datatypes avaialble filters"""

    for i, dataset in enumerate(datasets):
        slices = dataset.slice_orthogonal()
        assert slices is not None
        assert isinstance(slices, vtki.MultiBlock)
        assert slices.n_blocks == 3
        for slc in slices:
            assert isinstance(slc, vtki.PolyData)

def test_slice_along_axis():
    """Test the many slices along axis filter """
    axii = ['x', 'y', 'z', 'y', 'x']
    ns =  [2, 3, 4, 10, 20, 13]
    for i, dataset in enumerate(datasets):
        slices = dataset.slice_along_axis(n=ns[i], axis=axii[i])
        assert slices is not None
        assert isinstance(slices, vtki.MultiBlock)
        assert slices.n_blocks == ns[i]
        for slc in slices:
            assert isinstance(slc, vtki.PolyData)

def test_threshold():
    for i, dataset in enumerate(datasets[0:3]):
        thresh = dataset.threshold()
        assert thresh is not None
        assert isinstance(thresh, vtki.UnstructuredGrid)
    # Test value ranges
    dataset = examples.load_uniform() # UniformGrid
    thresh = dataset.threshold(100, invert=False)
    assert thresh is not None
    assert isinstance(thresh, vtki.UnstructuredGrid)
    thresh = dataset.threshold([100, 500], invert=False)
    assert thresh is not None
    assert isinstance(thresh, vtki.UnstructuredGrid)
    thresh = dataset.threshold([100, 500], invert=True)
    assert thresh is not None
    assert isinstance(thresh, vtki.UnstructuredGrid)
    # Now test datasets without arrays
    with pytest.raises(RuntimeError):
        for i, dataset in enumerate(datasets[3:-1]):
            thresh = dataset.threshold()
            assert thresh is not None
            assert isinstance(thresh, vtki.UnstructuredGrid)

def test_threshold_percent():
    percents = [25, 50, [18.0, 85.0], [19.0, 80.0], 0.70]
    inverts = [False, True, False, True, False]
    # Only test data sets that have arrays
    for i, dataset in enumerate(datasets[0:3]):
        thresh = dataset.threshold_percent(percent=percents[i], invert=inverts[i])
        assert thresh is not None
        assert isinstance(thresh, vtki.UnstructuredGrid)

def test_outline():
    for i, dataset in enumerate(datasets):
        outline = dataset.outline()
        assert outline is not None
        assert isinstance(outline, vtki.PolyData)

def test_outline_corners():
    for i, dataset in enumerate(datasets):
        outline = dataset.outline_corners()
        assert outline is not None
        assert isinstance(outline, vtki.PolyData)

def test_extract_geometry():
    for i, dataset in enumerate(datasets):
        outline = dataset.extract_geometry()
        assert outline is not None
        assert isinstance(outline, vtki.PolyData)

def test_extract_edges():
    for i, dataset in enumerate(datasets):
        wire = dataset.extract_geometry()
        assert wire is not None
        assert isinstance(wire, vtki.PolyData)
        wire = dataset.wireframe()
        assert wire is not None
        assert isinstance(wire, vtki.PolyData)

def test_contour():
    dataset = examples.load_uniform()
    iso = dataset.contour()
    assert iso is not None
    iso = dataset.contour(isosurfaces=[100, 300, 500])
    assert iso is not None
<commit_msg>Fix typo<commit_after>import pytest

import vtki
from vtki import examples

datasets = [
    examples.load_uniform(), # UniformGrid
    examples.load_rectilinear(), # RectilinearGrid
    examples.load_hexbeam(), # UnstructuredGrid
    examples.load_airplane(), # PolyData
    examples.load_structured(), # StructuredGrid
]
normals = ['x', 'y', '-z', (1,1,1), (3.3, 5.4, 0.8)]

def test_clip_filter():
    """This tests the clip filter on all datatypes avaialble filters"""
    for i, dataset in enumerate(datasets):
        clp = dataset.clip(normal=normals[i], invert=True)
        assert clp is not None
        assert isinstance(clp, vtki.UnstructuredGrid)

def test_slice_filter():
    """This tests the slice filter on all datatypes avaialble filters"""
    for i, dataset in enumerate(datasets):
        slc = dataset.slice(normal=normals[i])
        assert slc is not None
        assert isinstance(slc, vtki.PolyData)

def test_slice_orthogonal_filter():
    """This tests the slice filter on all datatypes avaialble filters"""

    for i, dataset in enumerate(datasets):
        slices = dataset.slice_orthogonal()
        assert slices is not None
        assert isinstance(slices, vtki.MultiBlock)
        assert slices.n_blocks == 3
        for slc in slices:
            assert isinstance(slc, vtki.PolyData)

def test_slice_along_axis():
    """Test the many slices along axis filter """
    axii = ['x', 'y', 'z', 'y', 'x']
    ns =  [2, 3, 4, 10, 20, 13]
    for i, dataset in enumerate(datasets):
        slices = dataset.slice_along_axis(n=ns[i], axis=axii[i])
        assert slices is not None
        assert isinstance(slices, vtki.MultiBlock)
        assert slices.n_blocks == ns[i]
        for slc in slices:
            assert isinstance(slc, vtki.PolyData)

def test_threshold():
    for i, dataset in enumerate(datasets[0:3]):
        thresh = dataset.threshold()
        assert thresh is not None
        assert isinstance(thresh, vtki.UnstructuredGrid)
    # Test value ranges
    dataset = examples.load_uniform() # UniformGrid
    thresh = dataset.threshold(100, invert=False)
    assert thresh is not None
    assert isinstance(thresh, vtki.UnstructuredGrid)
    thresh = dataset.threshold([100, 500], invert=False)
    assert thresh is not None
    assert isinstance(thresh, vtki.UnstructuredGrid)
    thresh = dataset.threshold([100, 500], invert=True)
    assert thresh is not None
    assert isinstance(thresh, vtki.UnstructuredGrid)
    # Now test datasets without arrays
    with pytest.raises(RuntimeError):
        for i, dataset in enumerate(datasets[3:-1]):
            thresh = dataset.threshold()
            assert thresh is not None
            assert isinstance(thresh, vtki.UnstructuredGrid)

def test_threshold_percent():
    percents = [25, 50, [18.0, 85.0], [19.0, 80.0], 0.70]
    inverts = [False, True, False, True, False]
    # Only test data sets that have arrays
    for i, dataset in enumerate(datasets[0:3]):
        thresh = dataset.threshold_percent(percent=percents[i], invert=inverts[i])
        assert thresh is not None
        assert isinstance(thresh, vtki.UnstructuredGrid)

def test_outline():
    for i, dataset in enumerate(datasets):
        outline = dataset.outline()
        assert outline is not None
        assert isinstance(outline, vtki.PolyData)

def test_outline_corners():
    for i, dataset in enumerate(datasets):
        outline = dataset.outline_corners()
        assert outline is not None
        assert isinstance(outline, vtki.PolyData)

def test_extract_geometry():
    for i, dataset in enumerate(datasets):
        outline = dataset.extract_geometry()
        assert outline is not None
        assert isinstance(outline, vtki.PolyData)

def test_extract_edges():
    for i, dataset in enumerate(datasets):
        wire = dataset.extract_geometry()
        assert wire is not None
        assert isinstance(wire, vtki.PolyData)
        wire = dataset.wireframe()
        assert wire is not None
        assert isinstance(wire, vtki.PolyData)

def test_contour():
    dataset = examples.load_uniform()
    iso = dataset.contour()
    assert iso is not None
    iso = dataset.contour(isosurfaces=[100, 300, 500])
    assert iso is not None
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Tests for folding models."""

import numpy as np
import pytest

from parameters import T_VALUES

@pytest.mark.parametrize("t_values", T_VALUES)
@pytest.mark.parametrize("supercell_size", [(1, 1, 1), (2, 1, 1), (2, 1, 3)])
def test_fold_supercell_simple(
    get_model, t_values, supercell_size, sparse, models_close
):
    """
    Test that creating a supercell model and then folding it back creates
    the same model.
    """
    model = get_model(
        *t_values,
        sparse=sparse,
        uc=[[1, 0, 0.5], [0.1, 0.4, 0.0], [0.0, 0.0, 1.2]],
        pos=[[0, 0, 0], [0.2, 0.3, 0.1]]
    )
    supercell_model = model.supercell(size=supercell_size)
    orbital_labels = ["a", "b"] * np.prod(supercell_size)
    for i, offset_red in enumerate(supercell_model.pos[::2]):
        # TODO: changing the 'offset_red' manually is a temporary fix:
        # to be removed when more complex orbital matching is implemented.
        offset_cart = supercell_model.uc.T @ (offset_red - 1e-12)
        with pytest.warns(UserWarning):
            folded_model = supercell_model.fold_model(
                new_unit_cell=model.uc,
                unit_cell_offset=offset_cart,
                orbital_labels=orbital_labels,
                target_indices=[2 * i, 2 * i + 1],
            )
        assert models_close(model, folded_model)

def test_fold_inexact_positions(get_model, models_close):
    """
    Test that creating a supercell model and then folding it back creates
    the same model.
    """
    model = get_model(0.1, 0.3)
    model.uc = np.array([[1, 0, 0.5], [0.1, 0.4, 0.0], [0.0, 0.0, 1.2]])
    model.pos = np.array([[0, 0, 0], [0.5, 0.5, 0.5]])
    supercell_model = model.supercell(size=(8, 1, 1))
    orbital_labels = ["a", "b"] * 8
    np.random.seed(42)
    for i in range(len(supercell_model.pos)):
        # do not move positions around "base" unit cell
        if i in range(3, 7):
            continue
        delta = np.random.uniform(-0.01, 0.01, 3)
        supercell_model.pos[i] += delta
    with pytest.warns(UserWarning):
        folded_model = supercell_model.fold_model(
            new_unit_cell=model.uc,
            unit_cell_offset=supercell_model.uc.T @ supercell_model.pos[4],
            orbital_labels=orbital_labels,
            position_tolerance=0.1,
            check_cc=False,
        )
    assert models_close(model, folded_model)

@pytest.fixture
def get_model_pos_outside(get_model):
    """
    Fixture which creates a model where one position is outside the
    unit cell.
    """

    def inner():
        model = get_model(0.1, 0.3, uc=np.eye(3))
        model.pos[0][0] = -0.01
        return model

    return inner

def test_consistency_checks_disabled(
    get_model_pos_outside,
):  # pylint: disable=redefined-outer-name,invalid-name
    """
    Test that no error is raised if the folding consistency checks
    are disabled.
    """
    model = get_model_pos_outside()
    model.fold_model(
        new_unit_cell=model.uc,
        orbital_labels=["a", "b"],
        check_orbital_ratio=False,
        check_uc_volume=False,
        unmatched_position_threshold=0.0,
    )

def test_orbital_number_consistency_check(
    get_model_pos_outside,
):  # pylint: disable=redefined-outer-name,invalid-name
    """
    Test that the orbital ratio check raises an error for a model
    that has a position outside the unit cell.
    """
    model = get_model_pos_outside()
    with pytest.raises(ValueError) as excinfo:
        model.fold_model(
            new_unit_cell=model.uc, orbital_labels=["a", "b"], check_uc_volume=False
        )
    assert "individual orbital numbers" in str(excinfo.value)

def test_unmatched_position_check(
    get_model_pos_outside,
):  # pylint: disable=redefined-outer-name
    """
    Test that the unmatched position check raises an error for
    a model that has a position outside the unit cell.
    """
    model = get_model_pos_outside()
    with pytest.raises(ValueError) as excinfo:
        model.fold_model(
            new_unit_cell=model.uc,
            orbital_labels=["a", "b"],
            check_orbital_ratio=False,
            check_uc_volume=False,
            unmatched_position_threshold=1.0,
        )
    assert "does not match any orbital in the new model" in str(excinfo.value)

def test_unmatched_position_check_wrong_uc(get_model):  # pylint: disable=invalid-name
    """
    Test that the unmatched position check raises an error when folding
    with a unit cell that does not match the current lattice.
    """
    model = get_model(t1=0.1, t2=0.6, uc=np.eye(3))
    supercell_model = model.supercell(size=(2, 2, 2))
    with pytest.raises(ValueError) as excinfo:
        supercell_model.fold_model(
            new_unit_cell=np.diag([1, 1, 1.2]),
            unit_cell_offset=(0.1, 0.1, 0.1),
            orbital_labels=["a", "b"] * 8,
            check_uc_volume=False,
        )
    assert "does not match any orbital in the new model" in str(excinfo.value)

def test_volume_consistency_check(
    get_model_pos_outside,
):  # pylint: disable=redefined-outer-name
    """
    Test that the unit cell volume check raises an error for a model
    that has a position outside the unit cell.
    """
    model = get_model_pos_outside()
    with pytest.raises(ValueError) as excinfo:
        model.fold_model(
            new_unit_cell=model.uc, orbital_labels=["a", "b"], check_orbital_ratio=False
        )
    assert "unit cell volume" in str(excinfo.value)

def test_fractional_occupation_consistency_check(
    get_model,
):  # pylint: disable=invalid-name
    """
    Check that an error is raised when the resulting number of
    occupations is fractional. This is tested by manually changing
    the occupation number in a supercell model
    """
    model = get_model(0.1, 0.3, uc=np.eye(3))
    model.pos += 0.01
    supercell_model = model.supercell(size=(2, 1, 1))
    supercell_model.occ += 1
    with pytest.raises(ValueError) as excinfo:
        supercell_model.fold_model(
            new_unit_cell=model.uc, orbital_labels=["a", "b"] * 2
        )
    assert "fractional" in str(excinfo.value)

def test_new_unit_cell_check(get_model):
    """
    Test the check that the new unit cell must lie within the current one.
    """
    model = get_model(0.1, 0.3, uc=np.diag([1, 2, 3]))
    supercell_model = model.supercell(size=(1, 2, 1))
    with pytest.raises(ValueError) as excinfo:
        supercell_model.fold_model(
            new_unit_cell=model.uc,
            unit_cell_offset=(0.2, 0, 0),
            orbital_labels=["a", "b"] * 2,
        )
    assert "new unit cell is not contained" in str(excinfo.value)
<commit_msg>Update folding tests to pass with the folding warnings.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Tests for folding models."""

import numpy as np
import pytest

from parameters import T_VALUES

@pytest.mark.parametrize("t_values", T_VALUES)
@pytest.mark.parametrize("supercell_size", [(1, 1, 1), (2, 1, 1), (2, 1, 3)])
def test_fold_supercell_simple(
    get_model, t_values, supercell_size, sparse, models_close
):
    """
    Test that creating a supercell model and then folding it back creates
    the same model.
    """
    model = get_model(
        *t_values,
        sparse=sparse,
        uc=[[1, 0, 0.5], [0.1, 0.4, 0.0], [0.0, 0.0, 1.2]],
        pos=[[0, 0, 0], [0.2, 0.3, 0.1]]
    )
    supercell_model = model.supercell(size=supercell_size)
    orbital_labels = ["a", "b"] * np.prod(supercell_size)
    for i, offset_red in enumerate(supercell_model.pos[::2]):
        # TODO: changing the 'offset_red' manually is a temporary fix:
        # to be removed when more complex orbital matching is implemented.
        offset_cart = supercell_model.uc.T @ (offset_red - 1e-12)
        with pytest.warns(UserWarning):
            folded_model = supercell_model.fold_model(
                new_unit_cell=model.uc,
                unit_cell_offset=offset_cart,
                orbital_labels=orbital_labels,
                target_indices=[2 * i, 2 * i + 1],
            )
        assert models_close(model, folded_model)

def test_fold_inexact_positions(get_model, models_close):
    """
    Test that creating a supercell model and then folding it back creates
    the same model.
    """
    model = get_model(0.1, 0.3)
    model.uc = np.array([[1, 0, 0.5], [0.1, 0.4, 0.0], [0.0, 0.0, 1.2]])
    model.pos = np.array([[0, 0, 0], [0.5, 0.5, 0.5]])
    supercell_model = model.supercell(size=(8, 1, 1))
    orbital_labels = ["a", "b"] * 8
    np.random.seed(42)
    for i in range(len(supercell_model.pos)):
        # do not move positions around "base" unit cell
        if i in range(3, 7):
            continue
        delta = np.random.uniform(-0.01, 0.01, 3)
        supercell_model.pos[i] += delta
    with pytest.warns(UserWarning):
        folded_model = supercell_model.fold_model(
            new_unit_cell=model.uc,
            unit_cell_offset=supercell_model.uc.T @ supercell_model.pos[4],
            orbital_labels=orbital_labels,
            position_tolerance=0.1,
            check_cc=False,
        )
    assert models_close(model, folded_model)

@pytest.fixture
def get_model_pos_outside(get_model):
    """
    Fixture which creates a model where one position is outside the
    unit cell.
    """

    def inner():
        model = get_model(0.1, 0.3, uc=np.eye(3))
        model.pos[0][0] = -0.01
        return model

    return inner

def test_consistency_checks_disabled(
    get_model_pos_outside,
):  # pylint: disable=redefined-outer-name,invalid-name
    """
    Test that no error is raised if the folding consistency checks
    are disabled.
    """
    model = get_model_pos_outside()
    with pytest.warns(UserWarning):
        model.fold_model(
            new_unit_cell=model.uc,
            orbital_labels=["a", "b"],
            check_orbital_ratio=False,
            check_uc_volume=False,
            unmatched_position_threshold=0.0,
        )

def test_orbital_number_consistency_check(
    get_model_pos_outside,
):  # pylint: disable=redefined-outer-name,invalid-name
    """
    Test that the orbital ratio check raises an error for a model
    that has a position outside the unit cell.
    """
    model = get_model_pos_outside()
    with pytest.raises(ValueError) as excinfo:
        with pytest.warns(UserWarning):
            model.fold_model(
                new_unit_cell=model.uc, orbital_labels=["a", "b"], check_uc_volume=False
            )
    assert "individual orbital numbers" in str(excinfo.value)

def test_unmatched_position_check(
    get_model_pos_outside,
):  # pylint: disable=redefined-outer-name
    """
    Test that the unmatched position check raises an error for
    a model that has a position outside the unit cell.
    """
    model = get_model_pos_outside()
    with pytest.raises(ValueError) as excinfo:
        with pytest.warns(UserWarning):
            model.fold_model(
                new_unit_cell=model.uc,
                orbital_labels=["a", "b"],
                check_orbital_ratio=False,
                check_uc_volume=False,
                unmatched_position_threshold=1.0,
            )
    assert "does not match any orbital in the new model" in str(excinfo.value)

def test_unmatched_position_check_wrong_uc(get_model):  # pylint: disable=invalid-name
    """
    Test that the unmatched position check raises an error when folding
    with a unit cell that does not match the current lattice.
    """
    model = get_model(t1=0.1, t2=0.6, uc=np.eye(3))
    supercell_model = model.supercell(size=(2, 2, 2))
    with pytest.raises(ValueError) as excinfo:
        supercell_model.fold_model(
            new_unit_cell=np.diag([1, 1, 1.2]),
            unit_cell_offset=(0.1, 0.1, 0.1),
            orbital_labels=["a", "b"] * 8,
            check_uc_volume=False,
        )
    assert "does not match any orbital in the new model" in str(excinfo.value)

def test_volume_consistency_check(
    get_model_pos_outside,
):  # pylint: disable=redefined-outer-name
    """
    Test that the unit cell volume check raises an error for a model
    that has a position outside the unit cell.
    """
    model = get_model_pos_outside()
    with pytest.raises(ValueError) as excinfo:
        with pytest.warns(UserWarning):
            model.fold_model(
                new_unit_cell=model.uc,
                orbital_labels=["a", "b"],
                check_orbital_ratio=False,
            )
    assert "unit cell volume" in str(excinfo.value)

def test_fractional_occupation_consistency_check(
    get_model,
):  # pylint: disable=invalid-name
    """
    Check that an error is raised when the resulting number of
    occupations is fractional. This is tested by manually changing
    the occupation number in a supercell model
    """
    model = get_model(0.1, 0.3, uc=np.eye(3))
    model.pos += 0.01
    supercell_model = model.supercell(size=(2, 1, 1))
    supercell_model.occ += 1
    with pytest.raises(ValueError) as excinfo:
        supercell_model.fold_model(
            new_unit_cell=model.uc, orbital_labels=["a", "b"] * 2
        )
    assert "fractional" in str(excinfo.value)

def test_new_unit_cell_check(get_model):
    """
    Test the check that the new unit cell must lie within the current one.
    """
    model = get_model(0.1, 0.3, uc=np.diag([1, 2, 3]))
    supercell_model = model.supercell(size=(1, 2, 1))
    with pytest.raises(ValueError) as excinfo:
        supercell_model.fold_model(
            new_unit_cell=model.uc,
            unit_cell_offset=(0.2, 0, 0),
            orbital_labels=["a", "b"] * 2,
        )
    assert "new unit cell is not contained" in str(excinfo.value)
<|endoftext|>


<commit_before>import json
import pytest

from saleor.product.models import Category

def assert_success_response(content):
    assert 'errors' not in content
    assert 'data' in content
    assert 'viewer' in content['data']

def assert_corresponding_categories(data, queryset):
    assert type(data) == list
    assert len(data) == queryset.count()
    for (category_data, category) in zip(data, queryset):
        assert category_data['name'] == category.name

@pytest.mark.django_db()
def test_category_query(client, product_in_stock):
    category = Category.objects.first()
    query = """
        query {
            viewer {
                category(pk: %(category_pk)s) {
                    pk
                    name
                    productsCount
                    ancestors { name }
                    children { name }
                    siblings { name }
                }
            }
        }
    """ % {'category_pk': category.pk}
    response = client.post('/graphql/', {'query': query})
    content = json.loads(response.content)
    assert_success_response(content)
    category_data = content['data']['viewer']['category']
    assert category_data is not None
    assert int(category_data['pk']) == category.pk
    assert category_data['name'] == category.name
    assert category_data['productsCount'] == category.products.count()
    assert_corresponding_categories(
        category_data['ancestors'], category.get_ancestors())
    assert_corresponding_categories(
        category_data['children'], category.get_children())
    assert_corresponding_categories(
        category_data['siblings'], category.get_siblings())

@pytest.mark.django_db()
def test_product_query(client, product_in_stock):
    category = Category.objects.first()
    product = category.products.first()
    query = """
        query {
            viewer {
                category(pk: %(category_pk)s) {
                    products {
                        edges {
                            node {
                                pk
                                name
                                url
                                imageUrl
                                images { url }
                                variants {
                                    name
                                    stockQuantity
                                }
                                availability {
                                    available,
                                    priceRange {
                                        minPrice {
                                            gross
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    """ % {'category_pk': category.pk}
    response = client.post('/graphql/', {'query': query})
    content = json.loads(response.content)
    assert_success_response(content)
    assert content['data']['viewer']['category'] is not None
    product_edges_data = content['data']['viewer']['category']['products']['edges']
    assert len(product_edges_data) == category.products.count()
    product_data = product_edges_data[0]['node']
    assert int(product_data['pk']) == product.pk
    assert product_data['name'] == product.name
    assert product_data['url'] == product.get_absolute_url()
<commit_msg>Decode JSON response before deserializing it<commit_after>import json
import pytest

from saleor.product.models import Category

def get_content(response):
    return json.loads(response.content.decode('utf8'))

def assert_success_response(content):
    assert 'errors' not in content
    assert 'data' in content
    assert 'viewer' in content['data']

def assert_corresponding_categories(data, queryset):
    assert type(data) == list
    assert len(data) == queryset.count()
    for (category_data, category) in zip(data, queryset):
        assert category_data['name'] == category.name

@pytest.mark.django_db()
def test_category_query(client, product_in_stock):
    category = Category.objects.first()
    query = """
        query {
            viewer {
                category(pk: %(category_pk)s) {
                    pk
                    name
                    productsCount
                    ancestors { name }
                    children { name }
                    siblings { name }
                }
            }
        }
    """ % {'category_pk': category.pk}
    response = client.post('/graphql/', {'query': query})
    content = get_content(response)
    assert_success_response(content)
    category_data = content['data']['viewer']['category']
    assert category_data is not None
    assert int(category_data['pk']) == category.pk
    assert category_data['name'] == category.name
    assert category_data['productsCount'] == category.products.count()
    assert_corresponding_categories(
        category_data['ancestors'], category.get_ancestors())
    assert_corresponding_categories(
        category_data['children'], category.get_children())
    assert_corresponding_categories(
        category_data['siblings'], category.get_siblings())

@pytest.mark.django_db()
def test_product_query(client, product_in_stock):
    category = Category.objects.first()
    product = category.products.first()
    query = """
        query {
            viewer {
                category(pk: %(category_pk)s) {
                    products {
                        edges {
                            node {
                                pk
                                name
                                url
                                imageUrl
                                images { url }
                                variants {
                                    name
                                    stockQuantity
                                }
                                availability {
                                    available,
                                    priceRange {
                                        minPrice {
                                            gross
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    """ % {'category_pk': category.pk}
    response = client.post('/graphql/', {'query': query})
    content = get_content(response)
    assert_success_response(content)
    assert content['data']['viewer']['category'] is not None
    product_edges_data = content['data']['viewer']['category']['products']['edges']
    assert len(product_edges_data) == category.products.count()
    product_data = product_edges_data[0]['node']
    assert int(product_data['pk']) == product.pk
    assert product_data['name'] == product.name
    assert product_data['url'] == product.get_absolute_url()
<|endoftext|>


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

from __future__ import print_function, division, unicode_literals, absolute_import

import datetime as dt

import numpy as np
from numpy.testing import assert_allclose
import pytest

from apexpy import helpers

##############################################################################
# NOTE: whenever function outputs are tested against hard-coded numbers,     #
# the test results (numbers) were obtained by running the code that is       #
# tested. Therefore these tests below only check that nothing changes when   #
# refactoring etc., and not if the results are actually correct              #
##############################################################################

###============================================================================
### Test constants
###============================================================================

def test_constants():
    assert_allclose(helpers.WGS84_e2, 0.006694379990141317)
    assert_allclose(helpers.WGS84_a, 6378.137)
    assert_allclose(helpers.d2r, np.pi/180)
    assert_allclose(helpers.r2d, 180/np.pi)

###============================================================================
### Test checklat
###============================================================================

def test_checklat_scalar():
    assert helpers.checklat(90) == 90
    assert helpers.checklat(0) == 0
    assert helpers.checklat(-90) == -90

    assert helpers.checklat(90+1e-5) == 90
    assert helpers.checklat(-90-1e-5) == -90

    assert type(helpers.checklat(0.)) == np.float64
    assert type(helpers.checklat(0)) == int
    assert type(helpers.checklat(90+1e-5)) == int

    with pytest.raises(ValueError):
        helpers.checklat(90+1e-4)
    with pytest.raises(ValueError):
        helpers.checklat(-90-1e-4)

def test_checklat_message():
    with pytest.raises(ValueError) as excinfo:
        helpers.checklat(100)
    assert str(excinfo.value).startswith('lat must be in')
    with pytest.raises(ValueError) as excinfo:
        helpers.checklat(100, name='glat')
    assert str(excinfo.value).startswith('glat')

def test_checklat_array():
    assert_allclose(helpers.checklat([-90-1e-5, -90, 0, 90, 90+1e-5]), np.array([-90, -90, 0, 90, 90]), rtol=0, atol=1e-8)

    assert type(helpers.checklat([0])) == np.ndarray

    with pytest.raises(ValueError):
        helpers.checklat([-90-1e-4, -90, 0, 90, 90+1e-5])

    with pytest.raises(ValueError):
        helpers.checklat([-90-1e-5, -90, 0, 90, 90+1e-4])

###============================================================================
### Test getsinIm
###============================================================================

def test_getsinIm_single():
    assert_allclose(helpers.getsinIm(60), 0.96076892283052284)
    assert_allclose(helpers.getsinIm(10), 0.33257924500670238)
    assert type(helpers.getsinIm(60)) != np.ndarray

def test_getsinIm_1Darray():
    assert_allclose(helpers.getsinIm([60, 10]), [0.96076892283052284, 0.33257924500670238])

def test_getsinIm_2Darray():
    assert_allclose(helpers.getsinIm([[60, 10], [60, 10]]), [[0.96076892283052284, 0.33257924500670238],
                                                             [0.96076892283052284, 0.33257924500670238]])

###============================================================================
### Test getcosIm
###============================================================================

def test_getcosIm_single():
    assert_allclose(helpers.getcosIm(60), 0.27735009811261463)
    assert_allclose(helpers.getcosIm(10), 0.94307531289434765)
    assert type(helpers.getcosIm(60)) != np.ndarray

def test_getcosIm_1Darray():
    assert_allclose(helpers.getcosIm([60, 10]), [0.27735009811261463, 0.94307531289434765])

def test_getcosIm_2Darray():
    assert_allclose(helpers.getcosIm([[60, 10], [60, 10]]), [[0.27735009811261463, 0.94307531289434765],
                                                             [0.27735009811261463, 0.94307531289434765]])

###============================================================================
### Test toYearFraction
###============================================================================

def test_toYearFraction():
    assert_allclose(helpers.toYearFraction(dt.datetime(2001, 1, 1, 0, 0, 0)), 2001)
    assert_allclose(helpers.toYearFraction(dt.datetime(2002, 1, 1, 0, 0, 0)), 2002)
    assert_allclose(helpers.toYearFraction(dt.datetime(2005, 2, 3, 4, 5, 6)), 2005.090877283105)
    assert_allclose(helpers.toYearFraction(dt.datetime(2005, 12, 11, 10, 9, 8)), 2005.943624682902)

###============================================================================
### Test gc2gdlat
###============================================================================

def test_gc2gdlat():
    assert_allclose(helpers.gc2gdlat(0), 0)
    assert_allclose(helpers.gc2gdlat(90), 90)
    assert_allclose(helpers.gc2gdlat(30), 30.166923849507356)
    assert_allclose(helpers.gc2gdlat(60), 60.166364190170931)

###============================================================================
### Test subsol
###============================================================================

def test_subsol():
    assert_allclose(helpers.subsol(dt.datetime(2005, 2, 3, 4, 5, 6)), (-16.505391672592904, 122.17768157084515))
    assert_allclose(helpers.subsol(dt.datetime(2010, 12, 11, 10, 9, 8)), (-23.001554595838947, 26.008999999955023))

if __name__ == '__main__':
    pytest.main()
<commit_msg>python 3 compatibility<commit_after># -*- coding: utf-8 -*-

from __future__ import print_function, division, unicode_literals, absolute_import

import datetime as dt

import numpy as np
from numpy.testing import assert_allclose
import pytest

from apexpy import helpers

##############################################################################
# NOTE: whenever function outputs are tested against hard-coded numbers,     #
# the test results (numbers) were obtained by running the code that is       #
# tested. Therefore these tests below only check that nothing changes when   #
# refactoring etc., and not if the results are actually correct              #
##############################################################################

###============================================================================
### Test constants
###============================================================================

def test_constants():
    assert_allclose(helpers.WGS84_e2, 0.006694379990141317)
    assert_allclose(helpers.WGS84_a, 6378.137)
    assert_allclose(helpers.d2r, np.pi/180)
    assert_allclose(helpers.r2d, 180/np.pi)

###============================================================================
### Test checklat
###============================================================================

def test_checklat_scalar():
    assert helpers.checklat(90) == 90
    assert helpers.checklat(0) == 0
    assert helpers.checklat(-90) == -90

    assert helpers.checklat(90+1e-5) == 90
    assert helpers.checklat(-90-1e-5) == -90

    assert type(helpers.checklat(0.)) == np.float64
    assert type(helpers.checklat(0)) == np.float64
    assert type(helpers.checklat(90+1e-5)) == np.float64

    with pytest.raises(ValueError):
        helpers.checklat(90+1e-4)
    with pytest.raises(ValueError):
        helpers.checklat(-90-1e-4)

def test_checklat_message():
    with pytest.raises(ValueError) as excinfo:
        helpers.checklat(100)
    assert str(excinfo.value).startswith('lat must be in')
    with pytest.raises(ValueError) as excinfo:
        helpers.checklat(100, name='glat')
    assert str(excinfo.value).startswith('glat')

def test_checklat_array():
    assert_allclose(helpers.checklat([-90-1e-5, -90, 0, 90, 90+1e-5]), np.array([-90, -90, 0, 90, 90]), rtol=0, atol=1e-8)

    assert type(helpers.checklat([0])) == np.ndarray

    with pytest.raises(ValueError):
        helpers.checklat([-90-1e-4, -90, 0, 90, 90+1e-5])

    with pytest.raises(ValueError):
        helpers.checklat([-90-1e-5, -90, 0, 90, 90+1e-4])

###============================================================================
### Test getsinIm
###============================================================================

def test_getsinIm_single():
    assert_allclose(helpers.getsinIm(60), 0.96076892283052284)
    assert_allclose(helpers.getsinIm(10), 0.33257924500670238)
    assert type(helpers.getsinIm(60)) != np.ndarray

def test_getsinIm_1Darray():
    assert_allclose(helpers.getsinIm([60, 10]), [0.96076892283052284, 0.33257924500670238])

def test_getsinIm_2Darray():
    assert_allclose(helpers.getsinIm([[60, 10], [60, 10]]), [[0.96076892283052284, 0.33257924500670238],
                                                             [0.96076892283052284, 0.33257924500670238]])

###============================================================================
### Test getcosIm
###============================================================================

def test_getcosIm_single():
    assert_allclose(helpers.getcosIm(60), 0.27735009811261463)
    assert_allclose(helpers.getcosIm(10), 0.94307531289434765)
    assert type(helpers.getcosIm(60)) != np.ndarray

def test_getcosIm_1Darray():
    assert_allclose(helpers.getcosIm([60, 10]), [0.27735009811261463, 0.94307531289434765])

def test_getcosIm_2Darray():
    assert_allclose(helpers.getcosIm([[60, 10], [60, 10]]), [[0.27735009811261463, 0.94307531289434765],
                                                             [0.27735009811261463, 0.94307531289434765]])

###============================================================================
### Test toYearFraction
###============================================================================

def test_toYearFraction():
    assert_allclose(helpers.toYearFraction(dt.datetime(2001, 1, 1, 0, 0, 0)), 2001)
    assert_allclose(helpers.toYearFraction(dt.datetime(2002, 1, 1, 0, 0, 0)), 2002)
    assert_allclose(helpers.toYearFraction(dt.datetime(2005, 2, 3, 4, 5, 6)), 2005.090877283105)
    assert_allclose(helpers.toYearFraction(dt.datetime(2005, 12, 11, 10, 9, 8)), 2005.943624682902)

###============================================================================
### Test gc2gdlat
###============================================================================

def test_gc2gdlat():
    assert_allclose(helpers.gc2gdlat(0), 0)
    assert_allclose(helpers.gc2gdlat(90), 90)
    assert_allclose(helpers.gc2gdlat(30), 30.166923849507356)
    assert_allclose(helpers.gc2gdlat(60), 60.166364190170931)

###============================================================================
### Test subsol
###============================================================================

def test_subsol():
    assert_allclose(helpers.subsol(dt.datetime(2005, 2, 3, 4, 5, 6)), (-16.505391672592904, 122.17768157084515))
    assert_allclose(helpers.subsol(dt.datetime(2010, 12, 11, 10, 9, 8)), (-23.001554595838947, 26.008999999955023))

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