<commit_before>#!/usr/bin/env python
#
# Code to interpret XDS (G)XPARM file and produce the rotation axis and beam
# vector in the imgCIF coordinate frame. From there, can transform to
# Mosflm missetting angles.
#

from scitbx import matrix
import math
import sys

def parse_xparm(xparm_file):
    '''Read an xparm file, return the rotation axis and beam vector in the
    XDS coordinate frame. Also the vector in this same coordinate frame to
    the start of the first pixel on the detector.'''

    values = map(float, open(xparm_file, 'r').read().split())

    ra = tuple(values[3:6])
    beam = tuple(values[7:10])

    # calculation of the true detector origin
    px, py = values[12], values[13]
    ox, oy = values[15], values[16]

    # question - what is the refined distance defined as *precisely* -
    # the distance from the crystal to the detector origin?! - if it's
    # mrad off, it will essentially make *no* difference

    x_to_d = - px * ox, - py * oy, values[14]

    return ra, beam, x_to_d, (px, py), values[14]

def xds_to_cbf(xparm_file):
    '''Given an XDS XPARM file, return a matrix which will transform from
    XDS coordinates to CBF reference frame.'''

    # XDS coordinate frame: x = detector fast direction
    #                       y = detector slow direction
    #                       z = x ^ y is detector normal (n.b. TOWARDS
    #                                                     the source)
    # CBF coordinate frame: x = primary gonio axis
    #                       y = z ^ x
    #                       z = component of beam |'r to x
    #
    # well that's ok then...

    ra, beam, x_to_d, pixel, distance = parse_xparm(xparm_file)

    # make them vectors

    ra = matrix.col(ra)
    beam = matrix.col(beam)
    x_to_d = matrix.col(x_to_d)

    if False:

        # try to reproduce the beam centre

        distance = x_to_d[2]
        nbeam = math.sqrt(beam.dot())

        bx_to_d = beam * distance / beam.elems[2]
        offset = bx_to_d - matrix.col(x_to_d)

        # ok this reproduces the correct beam centre - why can't I get this
        # below?!

        print 'Refined beam centre in pixels'
        print offset.elems[0] / pixel[0], offset.elems[1] / pixel[1]

    if False:

        # assert: distance is in direction of beam: convert to distance to
        # origin

        d1 = distance * beam / math.sqrt(beam.dot())
        d2 = d1.elems[2]

        x_to_d = x_to_d[0], x_to_d[1], d2

        print 'New distance: %.3f' % d2

    wavelength = 1.0 / math.sqrt(beam.dot())

    print 'Wavelength: %.5f' % wavelength

    # then unit vectors

    ra = ra / math.sqrt(ra.dot())
    beam = beam / math.sqrt(beam.dot())

    # right now lets calculate a rotation which will overlap the two reference
    # frames.... first relabel so we know what we are trying to do

    x = ra

    z = beam - (beam.dot(ra) * ra)

    z = z / math.sqrt(z.dot())

    y = z.cross(x)

    # now lets figure the rotations we want to rotate x to _x, etc.

    _x = matrix.col([1, 0, 0])
    _y = matrix.col([0, 1, 0])
    _z = matrix.col([0, 0, 1])

    # ok then - #1 rotate about x ^ (1, 0, 0) - if they are not identical

    if _x.dot(x):
        _ra_x = _x.cross(x)
        _a_x = _x.angle(x)
    else:
        _ra_x = _x
        _a_x = 0.0

    _m_x = _ra_x.axis_and_angle_as_r3_rotation_matrix(- _a_x)

    # then rotate z to _z by rotating about _x (which is now coincident
    # with x) N.B. z cannot be perpendicular to x

    _ra_z = _x

    _a_z = _z.angle(_m_x * z)

    _m_z = _ra_z.axis_and_angle_as_r3_rotation_matrix(- _a_z)

    # _m is matrix to rotate FROM xds coordinate frame TO cbf reference frame

    _m = _m_z * _m_x

    # verify that this is a rotation i.e. has determinant +1.

    assert(math.fabs(_m.determinant() - 1.0) < 1.0e-7)

    # now rotate the original things thus, to ensure that they are behaving
    # themselves.

    print '%10.7f %10.7f %10.7f' % (_m * x).elems
    print '%10.7f %10.7f %10.7f' % (_m * y).elems
    print '%10.7f %10.7f %10.7f' % (_m * z).elems

    # then see how things behave when we look at the rest of the environment

    print 'Following is in CBF reference frame'
    print 'New laboratory frame axes (detector axes, then normal)'

    print '%10.7f %10.7f %10.7f' % (_m * matrix.col([1, 0, 0])).elems
    print '%10.7f %10.7f %10.7f' % (_m * matrix.col([0, 1, 0])).elems
    print '%10.7f %10.7f %10.7f' % (_m * matrix.col([0, 0, 1])).elems

    # and finally the new beam

    print 'New beam vector'

    print '%10.7f %10.7f %10.7f' % (_m * beam).elems

    print 'To detector origin i.e. start of the first pixel'

    print '%10.4f %10.4f %10.4f' % (_m * x_to_d).elems

    # now need to consider the position of the detector etc. to derive the
    # new direct beam centre...

    # so that would be the distance between the sample and the
    # detector plane in the direction of the direct beam.

    n = _m * matrix.col([0, 0, 1])

    n = n / math.sqrt(n.dot())

    D = _m * x_to_d

    d = n.dot(D)

    # calculate beam unit vector in cbf coordinate frame
    b = _m * beam
    b = b / math.sqrt(b.dot())

    # this will be the exact position where the beam strikes the
    # detector face - in the CBF coordinate frame
    B = b * (d / (b.dot(n)))

    print '%10.4f %10.4f %10.4f' % B.elems

    # which I now need to convert to coordinates on the detector face!

    o = B - D

    # which is what w.r.t. the transformed x, y axes? - easy we know this
    # transformation - it's the inverse rotation!

    _o = _m.inverse() * o

    print '%10.4f %10.4f %10.4f' % _o.elems

    # check the result is in the detector plane
    assert(math.fabs(_o.elems[2]) < 1.0e-7)

    # right that's enough for today - though would be nice to plot the
    # refined beam coordinates.

    print _o.elems[0] / pixel[0], _o.elems[1] / pixel[1]

if __name__ == '__main__':
    xds_to_cbf(sys.argv[1])
<commit_msg>Think I have the conversion now to a refined mosflm coordinate frame, tidy<commit_after>#!/usr/bin/env python
#
# Code to interpret XDS (G)XPARM file and produce the rotation axis and beam
# vector in the imgCIF coordinate frame. From there, can transform to
# Mosflm missetting angles.
#

from scitbx import matrix
import math
import sys

def parse_xparm(xparm_file):
    '''Read an xparm file, return the rotation axis and beam vector in the
    XDS coordinate frame. Also the vector in this same coordinate frame to
    the start of the first pixel on the detector.'''

    values = map(float, open(xparm_file, 'r').read().split())

    ra = tuple(values[3:6])
    beam = tuple(values[7:10])

    # calculation of the true detector origin
    nx, ny = int(values[10]), int(values[11])
    px, py = values[12], values[13]
    ox, oy = values[15], values[16]

    # question - what is the refined distance defined as *precisely* -
    # the distance from the crystal to the detector origin?! - if it's
    # mrad off, it will essentially make *no* difference

    x_to_d = - px * ox, - py * oy, values[14]

    return ra, beam, x_to_d, (px, py), values[14], (nx, ny)

def xds_to_cbf(xparm_file):
    '''Given an XDS XPARM file, return a matrix which will transform from
    XDS coordinates to CBF reference frame.'''

    # XDS coordinate frame: x = detector fast direction
    #                       y = detector slow direction
    #                       z = x ^ y is detector normal (n.b. TOWARDS
    #                                                     the source)
    # CBF coordinate frame: x = primary gonio axis
    #                       y = z ^ x
    #                       z = component of beam |'r to x
    #
    # well that's ok then...

    ra, beam, x_to_d, pixel, distance, nxny = parse_xparm(xparm_file)

    # make them vectors

    ra = matrix.col(ra)
    beam = matrix.col(beam)
    x_to_d = matrix.col(x_to_d)
    nx, ny = nxny

    if False:

        # try to reproduce the beam centre

        distance = x_to_d[2]
        nbeam = math.sqrt(beam.dot())

        bx_to_d = beam * distance / beam.elems[2]
        offset = bx_to_d - matrix.col(x_to_d)

        # ok this reproduces the correct beam centre - why can't I get this
        # below?!

        print 'Refined beam centre in pixels'
        print offset.elems[0] / pixel[0], offset.elems[1] / pixel[1]

    if False:

        # assert: distance is in direction of beam: convert to distance to
        # origin

        d1 = distance * beam / math.sqrt(beam.dot())
        d2 = d1.elems[2]

        x_to_d = x_to_d[0], x_to_d[1], d2

        print 'New distance: %.3f' % d2

    wavelength = 1.0 / math.sqrt(beam.dot())

    print 'Wavelength: %.5f' % wavelength

    # then unit vectors

    ra = ra / math.sqrt(ra.dot())
    beam = beam / math.sqrt(beam.dot())

    # right now lets calculate a rotation which will overlap the two reference
    # frames.... first relabel so we know what we are trying to do

    x = ra

    z = beam - (beam.dot(ra) * ra)

    z = z / math.sqrt(z.dot())

    y = z.cross(x)

    # now lets figure the rotations we want to rotate x to _x, etc.

    _x = matrix.col([1, 0, 0])
    _y = matrix.col([0, 1, 0])
    _z = matrix.col([0, 0, 1])

    # ok then - #1 rotate about x ^ (1, 0, 0) - if they are not identical

    if _x.dot(x):
        _ra_x = _x.cross(x)
        _a_x = _x.angle(x)
    else:
        _ra_x = _x
        _a_x = 0.0

    _m_x = _ra_x.axis_and_angle_as_r3_rotation_matrix(- _a_x)

    # then rotate z to _z by rotating about _x (which is now coincident
    # with x) N.B. z cannot be perpendicular to x

    _ra_z = _x

    _a_z = _z.angle(_m_x * z)

    _m_z = _ra_z.axis_and_angle_as_r3_rotation_matrix(- _a_z)

    # _m is matrix to rotate FROM xds coordinate frame TO cbf reference frame

    _m = _m_z * _m_x

    # verify that this is a rotation i.e. has determinant +1.

    assert(math.fabs(_m.determinant() - 1.0) < 1.0e-7)

    # now rotate the original things thus, to ensure that they are behaving
    # themselves.

    print '%10.7f %10.7f %10.7f' % (_m * x).elems
    print '%10.7f %10.7f %10.7f' % (_m * y).elems
    print '%10.7f %10.7f %10.7f' % (_m * z).elems

    # then see how things behave when we look at the rest of the environment

    print 'Following is in CBF reference frame'
    print 'New laboratory frame axes (detector axes, then normal)'

    print '%10.7f %10.7f %10.7f' % (_m * matrix.col([1, 0, 0])).elems
    print '%10.7f %10.7f %10.7f' % (_m * matrix.col([0, 1, 0])).elems
    print '%10.7f %10.7f %10.7f' % (_m * matrix.col([0, 0, 1])).elems

    # and finally the new beam

    print 'New beam vector'

    print '%10.7f %10.7f %10.7f' % (_m * beam).elems

    print 'To detector origin i.e. start of the first pixel'

    print '%10.4f %10.4f %10.4f' % (_m * x_to_d).elems

    # now need to consider the position of the detector etc. to derive the
    # new direct beam centre...

    # so that would be the distance between the sample and the
    # detector plane in the direction of the direct beam.

    n = _m * matrix.col([0, 0, 1])

    n = n / math.sqrt(n.dot())

    D = _m * x_to_d

    d = n.dot(D)

    # calculate beam unit vector in cbf coordinate frame
    b = _m * beam
    b = b / math.sqrt(b.dot())

    # this will be the exact position where the beam strikes the
    # detector face - in the CBF coordinate frame
    B = b * (d / (b.dot(n)))

    print '%10.4f %10.4f %10.4f' % B.elems

    # which I now need to convert to coordinates on the detector face!

    o = B - D

    # which is what w.r.t. the transformed x, y axes? - easy we know this
    # transformation - it's the inverse rotation!

    _o = _m.inverse() * o

    print '%10.4f %10.4f %10.4f' % _o.elems

    # check the result is in the detector plane
    assert(math.fabs(_o.elems[2]) < 1.0e-7)

    # right that's enough for today - though would be nice to plot the
    # refined beam coordinates. OK that would be these then...

    print _o.elems[0] / pixel[0], _o.elems[1] / pixel[1]

    # now to calculate e.g. Mosflm corrections - TILT, TWIST, CCOMEGA
    # which are in 1/100 degrees for the first two, degrees for the
    # third.

    # to do this though we need to make the reference frame beam centric -
    # that is, rotate to ensure that the direct beam is along a cartesian
    # axis.

    # TILT - rotation about HORIZONTAL
    # TWIST - rotation about VERTICAL
    # CCOMEGA - rotation of detector about BEAM

    # first rotate about vertical axis

    _b = b
    _bt = matrix.col([0.0, 0.0, 1.0])
    _r = matrix.col([1.0, 0.0, 0.0])

    _v = _b.cross(_r)
    _a = _b.angle(_bt)

    cbf_to_mos = _v.axis_and_angle_as_r3_rotation_matrix(- _a)

    print '%10.7f %10.7f %10.7f' % _v.elems
    print 'angle: %f degrees' % (_a * 180.0 / math.pi)

    print '%10.7f %10.7f %10.7f' % (cbf_to_mos * _b).elems
    print '%10.7f %10.7f %10.7f' % (cbf_to_mos * _r).elems

    # now rotate everything else - in particular the detector origin, the
    # axes and the position on the detector where the direct beam hits
    # (which I assume is done in terms of pixel coordinates, converted
    # to mm...)

    # first attempt to get an accurate detector origin in the new (mosflm)
    # reference frame

    __o = cbf_to_mos * D

    print '%10.4f %10.4f %10.4f' % __o.elems

    __b = cbf_to_mos * B

    print '%10.4f %10.4f %10.4f' % __b.elems

    # get the new detector axes

    __x = cbf_to_mos * _m * matrix.col([1.0, 0.0, 0.0])
    __y = cbf_to_mos * _m * matrix.col([0.0, 1.0, 0.0])

    # now want to get the displacement from where the beam strikes the
    # detector to the origin w.r.t. these detector axes

    print '%10.7f %10.7f %10.7f' % __x.elems
    print '%10.7f %10.7f %10.7f' % __y.elems

    print 'Beam centre offset'

    __beam = cbf_to_mos * (B - D)

    print '%10.7f %10.7f %10.7f' % __beam.elems

    mos_to_det = (cbf_to_mos * _m * matrix.sqr([1.0, 0.0, 0.0,
                                                0.0, 1.0, 0.0,
                                                0.0, 0.0, 1.0])).inverse()

    print 'On detector coordinate frame'
    print '%10.7f %10.7f %10.7f' % (mos_to_det * __beam).elems

    x, y, z = (mos_to_det * __beam).elems

    print 'Mosflm beam centre'
    print '%10.4f %10.4f' % (nx * pixel[0] - x, ny * pixel[1] - y)

if __name__ == '__main__':
    xds_to_cbf(sys.argv[1])
<|endoftext|>


<commit_before>from compiler.ast import *
import os

import AST

# This is just and istance of Mul to use when we group numerators and
#  denominators
_EMPTY_MUL = Mul((None, None))

def dict2TeX(d, name_dict, lhs_form='%s', split_terms=False, simpleTeX=False):
    lines = []
    for lhs, rhs in d.items():
        if split_terms:
            ast = AST.strip_parse(rhs)
            pos, neg = [], []
            AST._collect_pos_neg(ast, pos, neg)
            try:
                lhsTeX = lhs_form % expr2TeX(lhs, name_dict=name_dict)
            except TypeError:
                lhsTeX = lhs_form
            rhsTeX = _ast2TeX(pos[0], name_dict=name_dict)
            lines.append(r'$ %s $ &=& $ %s $\\' % (lhsTeX, rhsTeX))

            for term in pos[1:]:
                TeXed = _ast2TeX(term, name_dict=name_dict)
                lines.append(r' & & $ + \, %s $\\' % TeXed)
            for term in neg:
                TeXed = _ast2TeX(term, name_dict=name_dict)
                lines.append(r' & & $ - \, %s $\\' % TeXed)
        else:
            lhsTeX = lhs_form % expr2TeX(lhs, name_dict=name_dict)
            rhsTeX = expr2TeX(rhs, name_dict=name_dict)
            lines.append(r'$ %s $ & = & $ %s $\\' % (lhsTeX, rhsTeX))

        if not simpleTeX:
            # Force a space between TeX'd entries
            lines[-1] = '%s[5mm]' % lines[-1]

    all = os.linesep.join(lines)

    if not simpleTeX:
        all = all.replace(r'\frac{', r'\tabfrac{')
        # This makes the fractions look much nicer in the tabular output. See
        #  http://www.texnik.de/table/table.phtml#fractions
        lines = [r'\providecommand{\tabfrac}[2]{%',
                 r'   \setlength{\fboxrule}{0pt}%',
                 r'   \fbox{$\frac{#1}{#2}$}}',
                 r'\begin{longtable}{rcl}'] + [all] + [r'\end{longtable}']
        all = os.linesep.join(lines)

    return all

def expr2TeX(expr, name_dict={}):
    """
    Return a TeX version of a python math expression.

    name_dict: A dictionary mapping variable names used in the expression to
        preferred TeX expressions.
    """
    ast = AST.strip_parse(expr)
    return _ast2TeX(ast, name_dict=name_dict)

def _ast2TeX(ast, outer=AST._FARTHEST_OUT, name_dict={},
             adjust=0):
    """
    Return a TeX version of an AST.

    outer: The AST's 'parent' node, used to determine whether or not to
        enclose the result in parentheses. The default of _FARTHEST_OUT will
        never enclose the result in parentheses.

    name_dict: A dictionary mapping variable names used in the expression to
        preferred TeX expressions.

    adjust: A numerical value to adjust the priority of this ast for
        particular cases. For example, the denominator of a '/' needs
        parentheses in more cases than does the numerator.
    """
    if isinstance(ast, Name):
        # Try to get a value from the name_dict, defaulting to ast.name if
        #  ast.name isn't in name_dict
        out = name_dict.get(ast.name, ast.name)
    elif isinstance(ast, Const):
        out = str(ast.value)
    elif isinstance(ast, Add):
        out = '%s + %s' % (_ast2TeX(ast.left, ast, name_dict),
                           _ast2TeX(ast.right, ast, name_dict))
    elif isinstance(ast, Sub):
        out = '%s - %s' % (_ast2TeX(ast.left, ast, name_dict),
                           _ast2TeX(ast.right, ast, name_dict, adjust = 1))
    elif isinstance(ast, Mul) or isinstance(ast, Div):
        # We collect all terms numerator and denominator
        nums, denoms = [], []
        AST._collect_num_denom(ast, nums, denoms)
        # _EMPTY_MUL ensures that parentheses are done properly, since every
        #  element is now the child of a Mul
        lam_func = lambda arg: _ast2TeX(arg, _EMPTY_MUL, name_dict)
        nums = [lam_func(term) for term in nums]
        if denoms:
            denoms = [lam_func(term) for term in denoms]
            out = r'\frac{%s}{%s}' % (r' \cdot '.join(nums),
                                      r' \cdot '.join(denoms))
        else:
            out = r' \cdot '.join(nums)
    elif isinstance(ast, Power):
        out = '{%s}^{%s}' % (_ast2TeX(ast.left, ast, name_dict, adjust = 1),
                             _ast2TeX(ast.right, ast, name_dict))
    elif isinstance(ast, UnarySub):
        out = '-%s' % _ast2TeX(ast.expr, ast, name_dict)
    elif isinstance(ast, UnaryAdd):
        out = '+%s' % _ast2TeX(ast.expr, ast, name_dict)
    elif isinstance(ast, CallFunc):
        lam_func = lambda arg: _ast2TeX(arg, name_dict=name_dict)
        name = lam_func(ast.node)
        args = [lam_func(arg) for arg in ast.args]
        if name == 'sqrt' and len(args) == 1:
            # Special case
            out = r'\sqrt{%s}' % args[0]
        else:
            out = r'\operatorname{%s}\left(%s\right)' % (name, r',\,'.join(args))

    if AST._need_parens(outer, ast, adjust):
        return out
    else:
        return r'\left(%s\right)' % out
<commit_msg>Tried to add support for and/or in py2TeX.<commit_after>from compiler.ast import *
import os

import AST

# This is just and istance of Mul to use when we group numerators and
#  denominators
_EMPTY_MUL = Mul((None, None))

def dict2TeX(d, name_dict, lhs_form='%s', split_terms=False, simpleTeX=False):
    lines = []
    for lhs, rhs in d.items():
        if split_terms:
            ast = AST.strip_parse(rhs)
            pos, neg = [], []
            AST._collect_pos_neg(ast, pos, neg)
            try:
                lhsTeX = lhs_form % expr2TeX(lhs, name_dict=name_dict)
            except TypeError:
                lhsTeX = lhs_form
            rhsTeX = _ast2TeX(pos[0], name_dict=name_dict)
            lines.append(r'$ %s $ &=& $ %s $\\' % (lhsTeX, rhsTeX))

            for term in pos[1:]:
                TeXed = _ast2TeX(term, name_dict=name_dict)
                lines.append(r' & & $ + \, %s $\\' % TeXed)
            for term in neg:
                TeXed = _ast2TeX(term, name_dict=name_dict)
                lines.append(r' & & $ - \, %s $\\' % TeXed)
        else:
            lhsTeX = lhs_form % expr2TeX(lhs, name_dict=name_dict)
            rhsTeX = expr2TeX(rhs, name_dict=name_dict)
            lines.append(r'$ %s $ & = & $ %s $\\' % (lhsTeX, rhsTeX))

        if not simpleTeX:
            # Force a space between TeX'd entries
            lines[-1] = '%s[5mm]' % lines[-1]

    all = os.linesep.join(lines)

    if not simpleTeX:
        all = all.replace(r'\frac{', r'\tabfrac{')
        # This makes the fractions look much nicer in the tabular output. See
        #  http://www.texnik.de/table/table.phtml#fractions
        lines = [r'\providecommand{\tabfrac}[2]{%',
                 r'   \setlength{\fboxrule}{0pt}%',
                 r'   \fbox{$\frac{#1}{#2}$}}',
                 r'\begin{longtable}{lll}'] + [all] + [r'\end{longtable}']
        all = os.linesep.join(lines)

    return all

def expr2TeX(expr, name_dict={}):
    """
    Return a TeX version of a python math expression.

    name_dict: A dictionary mapping variable names used in the expression to
        preferred TeX expressions.
    """
    ast = AST.strip_parse(expr)
    return _ast2TeX(ast, name_dict=name_dict)

def _ast2TeX(ast, outer=AST._FARTHEST_OUT, name_dict={},
             adjust=0):
    """
    Return a TeX version of an AST.

    outer: The AST's 'parent' node, used to determine whether or not to
        enclose the result in parentheses. The default of _FARTHEST_OUT will
        never enclose the result in parentheses.

    name_dict: A dictionary mapping variable names used in the expression to
        preferred TeX expressions.

    adjust: A numerical value to adjust the priority of this ast for
        particular cases. For example, the denominator of a '/' needs
        parentheses in more cases than does the numerator.
    """
    if isinstance(ast, Name):
        # Try to get a value from the name_dict, defaulting to ast.name if
        #  ast.name isn't in name_dict
        out = name_dict.get(ast.name, ast.name)
    elif isinstance(ast, Const):
        out = str(ast.value)
    elif isinstance(ast, Add):
        out = '%s + %s' % (_ast2TeX(ast.left, ast, name_dict),
                           _ast2TeX(ast.right, ast, name_dict))
    elif isinstance(ast, Sub):
        out = '%s - %s' % (_ast2TeX(ast.left, ast, name_dict),
                           _ast2TeX(ast.right, ast, name_dict, adjust = 1))
    elif isinstance(ast, Mul) or isinstance(ast, Div):
        # We collect all terms numerator and denominator
        nums, denoms = [], []
        AST._collect_num_denom(ast, nums, denoms)
        # _EMPTY_MUL ensures that parentheses are done properly, since every
        #  element is now the child of a Mul
        lam_func = lambda arg: _ast2TeX(arg, _EMPTY_MUL, name_dict)
        nums = [lam_func(term) for term in nums]
        if denoms:
            denoms = [lam_func(term) for term in denoms]
            out = r'\frac{%s}{%s}' % (r' \cdot '.join(nums),
                                      r' \cdot '.join(denoms))
        else:
            out = r' \cdot '.join(nums)
    elif isinstance(ast, Power):
        out = '{%s}^{%s}' % (_ast2TeX(ast.left, ast, name_dict, adjust = 1),
                             _ast2TeX(ast.right, ast, name_dict))
    elif isinstance(ast, UnarySub):
        out = '-%s' % _ast2TeX(ast.expr, ast, name_dict)
    elif isinstance(ast, UnaryAdd):
        out = '+%s' % _ast2TeX(ast.expr, ast, name_dict)
    elif isinstance(ast, CallFunc):
        lam_func = lambda arg: _ast2TeX(arg, name_dict=name_dict)
        name = lam_func(ast.node)
        args = [lam_func(arg) for arg in ast.args]
        if name == 'sqrt' and len(args) == 1:
            # Special case
            out = r'\sqrt{%s}' % args[0]
        else:
            out = r'\operatorname{%s}\left(%s\right)' % (name, r',\,'.join(args))
    elif isinstance(ast, Or) or isinstance(ast, And):
        out = r'\operatorname{%s}' % (str(ast))

    if AST._need_parens(outer, ast, adjust):
        return out
    else:
        return r'\left(%s\right)' % out
<|endoftext|>


<commit_before># coding=utf-8
import json

from django.contrib.auth.decorators import login_required
from django.http import HttpResponse
from django.shortcuts import render, get_object_or_404
from django.views.decorators.csrf import csrf_exempt

from frigg.helpers import github

from .models import Build, Project

@login_required
def overview(request):
    return render(request, "builds/overview.html", {
        'builds': Build.objects.all().order_by("-id").select_related('project', 'result')
    })

@login_required
def build(request, owner, name, build_number):
    return render(request, "builds/build.html", {
        'build': get_object_or_404(Build.objects.select_related('project'), project__owner=owner,
                                   project__name=name, build_number=build_number)
    })

@login_required
def deploy_master_branch(request, build_id):
    build = Build.objects.get(id=build_id)
    build.deploy()

    return HttpResponse("Deployed")

@csrf_exempt
def github_webhook(request):
    try:
        event = request.META['HTTP_X_GITHUB_EVENT']
    except KeyError:
        return HttpResponse("Missing HTTP_X_GITHUB_EVENT")

    data = json.loads(request.body)
    project = Project.objects.get_or_create_from_url(data['repo_url'])
    if event == "issue_comment":
        data = github.parse_comment_payload(data)

    elif event == "pull_request":
        data = github.parse_pull_request_payload(data)

    elif event == "push":
        data = github.parse_push_payload(data)

    else:
        return HttpResponse("Unknown event: %s" % event)

    if data:
        build = project.start_build(data)
        return HttpResponse('Handled "%s" event.\nMore info at %s' % (
            event,
            build.get_absolute_url())
        )
    else:
        return HttpResponse('Handled "%s" event.' % event)
<commit_msg>Move project creation<commit_after># coding=utf-8
import json

from django.contrib.auth.decorators import login_required
from django.http import HttpResponse
from django.shortcuts import render, get_object_or_404
from django.views.decorators.csrf import csrf_exempt

from frigg.helpers import github

from .models import Build, Project

@login_required
def overview(request):
    return render(request, "builds/overview.html", {
        'builds': Build.objects.all().order_by("-id").select_related('project', 'result')
    })

@login_required
def build(request, owner, name, build_number):
    return render(request, "builds/build.html", {
        'build': get_object_or_404(Build.objects.select_related('project'), project__owner=owner,
                                   project__name=name, build_number=build_number)
    })

@login_required
def deploy_master_branch(request, build_id):
    build = Build.objects.get(id=build_id)
    build.deploy()

    return HttpResponse("Deployed")

@csrf_exempt
def github_webhook(request):
    try:
        event = request.META['HTTP_X_GITHUB_EVENT']
    except KeyError:
        return HttpResponse("Missing HTTP_X_GITHUB_EVENT")

    data = json.loads(request.body)
    if event == "issue_comment":
        data = github.parse_comment_payload(data)

    elif event == "pull_request":
        data = github.parse_pull_request_payload(data)

    elif event == "push":
        data = github.parse_push_payload(data)

    else:
        return HttpResponse("Unknown event: %s" % event)

    if data:
        project = Project.objects.get_or_create_from_url(data['repo_url'])
        build = project.start_build(data)
        return HttpResponse('Handled "%s" event.\nMore info at %s' % (
            event,
            build.get_absolute_url())
        )
    else:
        return HttpResponse('Handled "%s" event.' % event)
<|endoftext|>


<commit_before>#!/usr/bin/python

'''
Calculating macroecological metrics for empirical or theoretical patch. Patch
is interpreted broadly as any temporally and spatially defined census.

Classes
-------
- `Patch` -- empirical metrics for census data

Patch Methods
-------------
- `sad` -- calculate species abundance distribution (grid or sample)
- `sar` -- calculate species-area relationship (grid or sample)
- `ear` -- calculate endemics-area relationship (grid or sample)
- `comm` -- calculate commonality between sub-patches (grid)
- `ssad` -- calculate species-level spatial abundance distrib (grid or sample)

- `get_sp_centers` --
- 'get_div_areas' -- return list of areas made by div_list

Misc functions
--------------
- `distance` -- return Euclidean distance between two points
'''

from __future__ import division
import numpy as np
from copy import deepcopy
from data import DataTable

__author__ = "Justin Kitzes"
__copyright__ = "Copyright 2012, Regents of the University of California"
__credits__ = ["John Harte"]
__license__ = None
__version__ = "0.1"
__maintainer__ = "Justin Kitzes"
__email__ = "jkitzes@berkeley.edu"
__status__ = "Development"

class Patch:
    '''
    An object representing an empirical census.

    Parameters
    ----------
    data_path : str
        Path to csv file containing census data.
    subset : dict
        Dictionary of permanent subset to data, {'column_name': 'condition'},
        which will limit all analysis to records in which column_name meets the
        condition, ie, {'year': '==2005', 'x': ('>20', '<40')} restricts
        analysis to year 2005 and x values between 20 and 40. These conditions
        can also be passed to the invididual methods, but subsetting the data
        table up front may save analysis time.

    Attributes
    ----------
    data_table : object of class DataTable
        Object containing patch data and metadata.

    '''

    def __init__(self, datapath, subset = {}):
        '''Initialize object of class Patch. See class documentation.'''

        self.data_table = DataTable(datapath)
        self.data_table.table = self.data_table.get_subtable(subset)

    def sad(self, criteria):
        '''
        Calculates an empirical species abundance distribution given criteria.

        Parameters
        ----------
        criteria : dict
            Dictionary of form {column_name: value}. Must contain a key with a
            value of 'species' indicating the column with species identifiers
            (this column must be type categorical in metadata). If a column
            giving the counts of species found at a point is also in the data,
            a key with the value 'count' should also be given.

            Value has a different meaning depending on column type:
            - metric - number of divisions of data along this axis, int/float
            - categorical - 'split' calculates each category separately

        Returns
        -------
        spp_list : ndarray
            1D array listing identifiers for species in the same order as they
            appear in arrays found in result.
        result : dict
            List of tuples containing results, where first element is
            dictionary of criteria for this calculation and second element is a
            1D ndarray of length species containing the abundance for each
            species.
        '''

        spp_list, spp_col, count_col, combinations = \
            self.parse_criteria(criteria)

        result = []
        for comb in combinations:

            subtable = self.data_table.get_subtable(comb)

            sad_list = []
            for species in spp_list:
                spp_subtable = subtable[subtable[spp_col] == species]
                if count_col:
                    count = np.sum(spp_subtable[count_col])
                else:
                    count = len(spp_subtable)
                sad_list.append(count)

            result.append((comb, np.array(sad_list)))

        # TODO: Consider changing array to structured with spp name as field

        return spp_list, result

    def parse_criteria(self, criteria):
        '''
        Parses criteria list to get all possible column combinations.

        Parameters
        ----------
        criteria : dict
            (See docstring for EPatch.sad)

        Returns
        -------
        spp_list : ndarray
            1D array listing identifiers for species in the same order as they
            appear in arrays found in result.
        spp_col : str
            Name of column containing species identifiers.
        count_col : str
            Name of column containing counts, if any.
        combinations : list of dicts
            List of dictionaries giving all possible combinations of criteria.
            Columns not mentioned in criteria are ignored and will be averaged
            over in later analyses.

        '''

        spp_list = None
        spp_col = None
        count_col = None
        combinations = []

        # Calculate all possible combinations of columns based on criteria
        # TODO: Add error checking
        for key, value in criteria.items():

            # Look for two special values indicating species and count cols
            if value == 'species':
                spp_list = np.unique(self.data_table.table[key])
                spp_col = key
                continue
            if value == 'count':
                count_col = key
                continue

            # Get levels of categorial or metric data
            if value == 'split':  # Categorial
                levels = np.unique(self.data_table.table[key])
                levels_str = ['==' + str(x) for x in levels]

            else:  # Metric

                # TODO: Make sure divisions equal if count (like LBRI), if true
                # continuous (like BCI) let divisions be anything. Currently no
                # checking.
                # TODO: Add support for sampling, if 'sample' in args, randomly
                # choose n start coords and store in levels

                dmin = self.data_table.meta[(key, 'minimum')]
                dmax = self.data_table.meta[(key, 'maximum')]
                dprec = self.data_table.meta[(key, 'precision')]

                # TODO: Error if step < prec
                step = (dmax + dprec - dmin) / value
                starts = np.arange(dmin, dmax + dprec, step)
                ends = starts + step

                starts_str = ['>=' + str(x) for x in starts]
                ends_str = ['<' + str(x) for x in ends]
                levels_str = zip(starts_str, ends_str)

            # Add these levels to combinations dictionary
            if len(combinations) == 0:  # If first criteria
                for i, level in enumerate(levels_str):
                    combinations.append({key: level})
            else:
                temp_comb = []
                for i, level in enumerate(levels_str):
                    exist_recs = deepcopy(combinations)
                    for rec in exist_recs:
                        rec[key] = level
                    temp_comb += exist_recs
                combinations = temp_comb

        return spp_list, spp_col, count_col, combinations

    def sar(self, div_cols, div_list, criteria, form='sar'):
        '''
        Calulate an empirical species-area relationship given criteria.

        Parameters
        ----------
        div_cols : tuple
            Column names to divide, eg, ('x', 'y'). Must be metric.
        div_list : list of tuples
            List of division pairs in same order as div_cols, eg, [(2,2),
            (2,4), (4,4)]. Values are number of divisions of div_col.
        criteria : dict
            See docstring for EPatch.sad. Here, criteria SHOULD NOT include
            items referring to div_cols (if there are any, they are ignored).
        form : string
            'sar' or 'ear' for species or endemics area relationship. EAR is
            relative to the subtable selected after criteria is applied.

        Returns
        -------
        areas : ndarray
            1D array of areas associated with each element of div_list.
        mean_result : ndarray
            1D array of same length as areas giving mean number of species or
            endemics at this area.
        full_result : list of ndarrays
            List of same length as areas containing arrays with element for
            count of species or endemics in each subpatch at corresponding
            area.
        '''

        # If any element in div_cols in criteria, remove from criteria
        criteria = {k: v for k, v in criteria.items() if k not in div_cols}

        # Loop through div combinations (ie, areas), calc sad, and summarize
        areas = []
        mean_result = []
        full_result = []

        for div in div_list:

            # Add divs to criteria dict
            this_criteria = deepcopy(criteria)
            for i, col in enumerate(div_cols):
                this_criteria[col] = div[i]

            # Get flattened sad for all criteria and this div
            spp_list, sad = self.sad(this_criteria)
            flat_sad = flatten_sad(sad)

            # Store results
            if type == 'sar':
                this_full = np.sum((flat_sad > 0), axis=0)
                this_mean = np.mean(this_full)
            elif type == 'ear':
                totcnt = np.sum(flat_sad, axis=1)
                totcnt_arr = \
                    np.array([list(totcnt),]*np.shape(flat_sad)[1]).transpose()

                this_full = np.sum(np.equal(flat_sad, totcnt_arr), axis=0)
                this_mean = np.mean(this_full)
            else:
                raise NotImplementedError('No SAR of type %s available' % type)

            full_result.append(this_full)
            mean_result.append(this_mean)

            # Store area
            area = 1
            for i, col in enumerate(div_cols):
                dmin = self.data_table.meta[(col, 'minimum')]
                dmax = self.data_table.meta[(col, 'maximum')]
                dprec = self.data_table.meta[(col, 'precision')]
                length = (dmax + dprec - dmin)

                area *= length / div[i]

            areas.append(area)

        # Return
        return areas, mean_result, full_result

def flatten_sad(sad):
    '''
    Takes a list of tuples, like sad output, ignores keys, and converts values
    into a 2D array with each value as a column (ie, species in rows, samples
    in columns.
    '''

    result = np.zeros((len(sad[0][1]), len(sad)))

    for i, tup in enumerate(sad):
        result[:,i] = tup[1]

    return result

def distance(pt1, pt2):
    ''' Calculate Euclidean distance between two points '''
    return np.sqrt((pt1[0] - pt2[0]) ** 2 + (pt1[1] - pt2[1]) ** 2)

def divisible(dividend, precision, divisor, tol = 1e-9):
    '''
    Check if dividend (here width or height of patch) is evenly divisible by
    divisor (here a number of patch divs) while accounting for floating point
    rounding issues.
    '''
    if divisor == 0:
        return False
    if divisor > round(dividend / precision):
        return False

    quot_raw = (dividend / precision) / divisor
    quot_round = round(quot_raw)
    diff = abs(quot_raw - quot_round)

    if diff < tol:
        return True
    else:
        return False

def rnd(num):
    '''
    Round num to number of decimal places in precision. Used to avoid issues
    with floating points in the patch and subpatch width and height that make
    subpatches not lie exactly on even divisions of patch.
    '''
    return round(num, 6)
<commit_msg>Correct sar arg to form throughout (replaces type)<commit_after>#!/usr/bin/python

'''
Calculating macroecological metrics for empirical or theoretical patch. Patch
is interpreted broadly as any temporally and spatially defined census.

Classes
-------
- `Patch` -- empirical metrics for census data

Patch Methods
-------------
- `sad` -- calculate species abundance distribution (grid or sample)
- `sar` -- calculate species-area relationship (grid or sample)
- `ear` -- calculate endemics-area relationship (grid or sample)
- `comm` -- calculate commonality between sub-patches (grid)
- `ssad` -- calculate species-level spatial abundance distrib (grid or sample)

- `get_sp_centers` --
- 'get_div_areas' -- return list of areas made by div_list

Misc functions
--------------
- `distance` -- return Euclidean distance between two points
'''

from __future__ import division
import numpy as np
from copy import deepcopy
from data import DataTable

__author__ = "Justin Kitzes"
__copyright__ = "Copyright 2012, Regents of the University of California"
__credits__ = ["John Harte"]
__license__ = None
__version__ = "0.1"
__maintainer__ = "Justin Kitzes"
__email__ = "jkitzes@berkeley.edu"
__status__ = "Development"

class Patch:
    '''
    An object representing an empirical census.

    Parameters
    ----------
    data_path : str
        Path to csv file containing census data.
    subset : dict
        Dictionary of permanent subset to data, {'column_name': 'condition'},
        which will limit all analysis to records in which column_name meets the
        condition, ie, {'year': '==2005', 'x': ('>20', '<40')} restricts
        analysis to year 2005 and x values between 20 and 40. These conditions
        can also be passed to the invididual methods, but subsetting the data
        table up front may save analysis time.

    Attributes
    ----------
    data_table : object of class DataTable
        Object containing patch data and metadata.

    '''

    def __init__(self, datapath, subset = {}):
        '''Initialize object of class Patch. See class documentation.'''

        self.data_table = DataTable(datapath)
        self.data_table.table = self.data_table.get_subtable(subset)

    def sad(self, criteria):
        '''
        Calculates an empirical species abundance distribution given criteria.

        Parameters
        ----------
        criteria : dict
            Dictionary of form {column_name: value}. Must contain a key with a
            value of 'species' indicating the column with species identifiers
            (this column must be type categorical in metadata). If a column
            giving the counts of species found at a point is also in the data,
            a key with the value 'count' should also be given.

            Value has a different meaning depending on column type:
            - metric - number of divisions of data along this axis, int/float
            - categorical - 'split' calculates each category separately

        Returns
        -------
        spp_list : ndarray
            1D array listing identifiers for species in the same order as they
            appear in arrays found in result.
        result : dict
            List of tuples containing results, where first element is
            dictionary of criteria for this calculation and second element is a
            1D ndarray of length species containing the abundance for each
            species.
        '''

        spp_list, spp_col, count_col, combinations = \
            self.parse_criteria(criteria)

        result = []
        for comb in combinations:

            subtable = self.data_table.get_subtable(comb)

            sad_list = []
            for species in spp_list:
                spp_subtable = subtable[subtable[spp_col] == species]
                if count_col:
                    count = np.sum(spp_subtable[count_col])
                else:
                    count = len(spp_subtable)
                sad_list.append(count)

            result.append((comb, np.array(sad_list)))

        # TODO: Consider changing array to structured with spp name as field

        return spp_list, result

    def parse_criteria(self, criteria):
        '''
        Parses criteria list to get all possible column combinations.

        Parameters
        ----------
        criteria : dict
            (See docstring for EPatch.sad)

        Returns
        -------
        spp_list : ndarray
            1D array listing identifiers for species in the same order as they
            appear in arrays found in result.
        spp_col : str
            Name of column containing species identifiers.
        count_col : str
            Name of column containing counts, if any.
        combinations : list of dicts
            List of dictionaries giving all possible combinations of criteria.
            Columns not mentioned in criteria are ignored and will be averaged
            over in later analyses.

        '''

        spp_list = None
        spp_col = None
        count_col = None
        combinations = []

        # Calculate all possible combinations of columns based on criteria
        # TODO: Add error checking
        for key, value in criteria.items():

            # Look for two special values indicating species and count cols
            if value == 'species':
                spp_list = np.unique(self.data_table.table[key])
                spp_col = key
                continue
            if value == 'count':
                count_col = key
                continue

            # Get levels of categorial or metric data
            if value == 'split':  # Categorial
                levels = np.unique(self.data_table.table[key])
                levels_str = ['==' + str(x) for x in levels]

            else:  # Metric

                # TODO: Make sure divisions equal if count (like LBRI), if true
                # continuous (like BCI) let divisions be anything. Currently no
                # checking.
                # TODO: Add support for sampling, if 'sample' in args, randomly
                # choose n start coords and store in levels

                dmin = self.data_table.meta[(key, 'minimum')]
                dmax = self.data_table.meta[(key, 'maximum')]
                dprec = self.data_table.meta[(key, 'precision')]

                # TODO: Error if step < prec
                step = (dmax + dprec - dmin) / value
                starts = np.arange(dmin, dmax + dprec, step)
                ends = starts + step

                starts_str = ['>=' + str(x) for x in starts]
                ends_str = ['<' + str(x) for x in ends]
                levels_str = zip(starts_str, ends_str)

            # Add these levels to combinations dictionary
            if len(combinations) == 0:  # If first criteria
                for i, level in enumerate(levels_str):
                    combinations.append({key: level})
            else:
                temp_comb = []
                for i, level in enumerate(levels_str):
                    exist_recs = deepcopy(combinations)
                    for rec in exist_recs:
                        rec[key] = level
                    temp_comb += exist_recs
                combinations = temp_comb

        return spp_list, spp_col, count_col, combinations

    def sar(self, div_cols, div_list, criteria, form='sar'):
        '''
        Calulate an empirical species-area relationship given criteria.

        Parameters
        ----------
        div_cols : tuple
            Column names to divide, eg, ('x', 'y'). Must be metric.
        div_list : list of tuples
            List of division pairs in same order as div_cols, eg, [(2,2),
            (2,4), (4,4)]. Values are number of divisions of div_col.
        criteria : dict
            See docstring for EPatch.sad. Here, criteria SHOULD NOT include
            items referring to div_cols (if there are any, they are ignored).
        form : string
            'sar' or 'ear' for species or endemics area relationship. EAR is
            relative to the subtable selected after criteria is applied.

        Returns
        -------
        areas : ndarray
            1D array of areas associated with each element of div_list.
        mean_result : ndarray
            1D array of same length as areas giving mean number of species or
            endemics at this area.
        full_result : list of ndarrays
            List of same length as areas containing arrays with element for
            count of species or endemics in each subpatch at corresponding
            area.
        '''

        # If any element in div_cols in criteria, remove from criteria
        criteria = {k: v for k, v in criteria.items() if k not in div_cols}

        # Loop through div combinations (ie, areas), calc sad, and summarize
        areas = []
        mean_result = []
        full_result = []

        for div in div_list:

            # Add divs to criteria dict
            this_criteria = deepcopy(criteria)
            for i, col in enumerate(div_cols):
                this_criteria[col] = div[i]

            # Get flattened sad for all criteria and this div
            spp_list, sad = self.sad(this_criteria)
            flat_sad = flatten_sad(sad)

            # Store results
            if form == 'sar':
                this_full = np.sum((flat_sad > 0), axis=0)
                this_mean = np.mean(this_full)
            elif form == 'ear':
                totcnt = np.sum(flat_sad, axis=1)
                totcnt_arr = \
                    np.array([list(totcnt),]*np.shape(flat_sad)[1]).transpose()

                this_full = np.sum(np.equal(flat_sad, totcnt_arr), axis=0)
                this_mean = np.mean(this_full)
            else:
                raise NotImplementedError('No SAR of form %s available' % form)

            full_result.append(this_full)
            mean_result.append(this_mean)

            # Store area
            area = 1
            for i, col in enumerate(div_cols):
                dmin = self.data_table.meta[(col, 'minimum')]
                dmax = self.data_table.meta[(col, 'maximum')]
                dprec = self.data_table.meta[(col, 'precision')]
                length = (dmax + dprec - dmin)

                area *= length / div[i]

            areas.append(area)

        # Return
        return areas, mean_result, full_result

def flatten_sad(sad):
    '''
    Takes a list of tuples, like sad output, ignores keys, and converts values
    into a 2D array with each value as a column (ie, species in rows, samples
    in columns.
    '''

    result = np.zeros((len(sad[0][1]), len(sad)))

    for i, tup in enumerate(sad):
        result[:,i] = tup[1]

    return result

def distance(pt1, pt2):
    ''' Calculate Euclidean distance between two points '''
    return np.sqrt((pt1[0] - pt2[0]) ** 2 + (pt1[1] - pt2[1]) ** 2)

def divisible(dividend, precision, divisor, tol = 1e-9):
    '''
    Check if dividend (here width or height of patch) is evenly divisible by
    divisor (here a number of patch divs) while accounting for floating point
    rounding issues.
    '''
    if divisor == 0:
        return False
    if divisor > round(dividend / precision):
        return False

    quot_raw = (dividend / precision) / divisor
    quot_round = round(quot_raw)
    diff = abs(quot_raw - quot_round)

    if diff < tol:
        return True
    else:
        return False

def rnd(num):
    '''
    Round num to number of decimal places in precision. Used to avoid issues
    with floating points in the patch and subpatch width and height that make
    subpatches not lie exactly on even divisions of patch.
    '''
    return round(num, 6)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Copyright (c) 2007-2008 rPath, Inc.
#
# This program is distributed under the terms of the Common Public License,
# version 1.0. A copy of this license should have been distributed with this
# source file in a file called LICENSE. If it is not present, the license
# is always available at http://www.rpath.com/permanent/licenses/CPL-1.0.
#
# This program is distributed in the hope that it will be useful, but
# without any warranty; without even the implied warranty of merchantability
# or fitness for a particular purpose. See the Common Public License for
# full details.
#

# Cristian Gafton, <gafton@rpath.com>
# handles extracting and inserting into various backends for database migration

import sys

from conary import dbstore

from schema import getTables, getIndexes
from tablelist import TableList

# database loader that knows how to insert rows efficiently
class Loader:
    binaryCols = [
        "salt", "password", "data", "entitlement", "pgpkey",
        "stream", "fileid", "pathid", "sha1" ]
    def __init__(self, db, table, fields):
        self.db = db
        self.table = table
        self.sql = "INSERT INTO %s (%s) VALUES (%s)" % (
            table, ", ".join(fields), ",".join(["?"]*len(fields)))
        self.cu = self.db.cursor()
        # check the fields are all present in the target
        self.cu.execute("SELECT * FROM %s LIMIT 1" % (table,))
        set1 = set([x.lower() for x in fields])
        set2 = set([x.lower() for x in self.cu.fields()])
        assert (set1.difference(set2) == set()), "Fields differ from src to target: %s %s" %(set1, set2)
        self.fields = fields
        # make sure we know when to use cu.binary()
        self.funcs = [ self.__getfunc(x) for x in fields ]
    def announce(self):
        sys.stdout.write("1by1 load table: %s\r" % (self.table,))
        sys.stdout.flush()
    # fields which we should tag as binary
    def __getfunc(self, field):
        if field.lower() in self.binaryCols:
            return self.cu.binary
        return lambda a: a
    # fix up the values in a row to match the transforms
    def rowvals(self, row):
        return tuple(map(lambda (f, x): f(x), zip(self.funcs, row)))
    def insertRows(self, rows, callback=None):
        self.cu.executemany(self.sql, [self.rowvals(row) for row in rows])
        if callback:
            callback.increment(len(rows))
        return len(rows)

class PgSQLLoader(Loader):
    def __init__(self, db, table, fields):
        Loader.__init__(self, db, table, fields)
        self.__usebulk = False
        if hasattr(self.db.dbh, "bulkload"):
            self.__usebulk = True
        if not self.__usebulk:
            print "WARNING: not using bulk load, update your python-pgsql bindings!"
    def bulkInsert(self, rows, callback):
        self.db.dbh.bulkload(self.table, (self.rowvals(row) for row in rows), self.fields)
        if callback:
            callback.increment(len(rows))
        return len(rows)
    def insertRows(self, rows, callback=None):
        if self.__usebulk:
            return self.bulkInsert(rows, callback)
        else:
            return Loader.insertRows(self, rows, callback)

class Database:
    def __init__(self, driver, db, verbose=True):
        self.db = dbstore.connect(db, driver)
        self.driver = driver
        self.db.loadSchema()
        self.verbose = verbose
    def createSchema(self):
        # create the tables, avoid the indexes
        cu = self.db.cursor()
        for stmt in getTables(self.driver):
            if self.verbose:
                print stmt
            cu.execute(stmt)
        self.db.loadSchema()
    def createIndexes(self):
        cu = self.db.cursor()
        for stmt in getIndexes(self.driver):
            if self.verbose:
                print stmt
            cu.execute(stmt)
        self.db.loadSchema()
    # check self.db.tables against the TableList
    def checkTablesList(self, isSrc=True):
        #  check that we are migrating all the tables in the source
        self.db.loadSchema()
        skip = ['databaseversion', 'instructionsets', 'commitlock']
        knowns = [x.lower() for x in TableList]
        haves = [x.lower() for x in self.db.tables]
        if isSrc:
            which = "Source"
        else:
            which = "Target"
        # tableList should not have items not present in the db
        onlyKnowns = set(knowns).difference(set(haves)).difference(set(skip))
        if onlyKnowns:
            raise RuntimeError("%s schema (%s) does not have table(s) %s" %(
                which, self.driver, onlyKnowns))
        # we should not have extra tables in the source
        onlyHaves = set(haves).difference(set(knowns)).difference(set(skip))
        if onlyHaves and isSrc:
            raise RuntimeError("TableList needs to be updated to handle tables", onlyHaves)
        return True

    # functions for when the instance is a source
    def getCount(self, table):
        cu = self.db.cursor()
        cu.execute("select count(*) from %s" % (table,))
        return cu.fetchall()[0][0]
    def getFields(self, table):
        cu = self.db.cursor()
        cu.execute("SELECT * FROM %s LIMIT 1" % (table,))
        return [x.lower() for x in cu.fields()]
    def getTables(self):
        return [x.lower() for x in self.db.tables]
    def iterRows(self, table, fields = "*"):
        cu = self.db.itercursor()
        cu.execute("select %s from %s" % (fields, table))
        return cu
    # functions for when the instance is a target
    def prepareInsert(self, table, fields):
        return Loader(self.db, table, fields)
    def finalize(self, version):
        self.db.setVersion(version)
        self.db.commit()
    # useful shortcut
    def commit(self):
        return self.db.commit()
    def close(self):
        self.db.close()

class PgSQLDatabase(Database):
    def __init__(self, db, verbose=True):
        Database.__init__(self, "postgresql", db, verbose)
    def createIndexes(self):
        Database.createIndexes(self)
        # fix the primary keys
        self.fix_primary_keys()
        # update the primary key sequences for all tables
    def fix_primary_keys(self):
        cu = self.db.cursor()
        # get the name of the primary key
        cu.execute("""
        select
            t.relname as table_name,
            col.attname as column_name
        from pg_class t
        join pg_index i on t.oid = i.indrelid and i.indisprimary = true
        join pg_class ind on i.indexrelid = ind.oid
        join pg_attribute col on col.attrelid = t.oid and col.attnum = i.indkey[0]
        where i.indnatts = 1
          and pg_catalog.pg_table_is_visible(t.oid)
        """)
        for (t, col) in cu.fetchall():
            table = t.lower()
            cu.execute("select pg_catalog.pg_get_serial_sequence(?, ?)", (table, col))
            seqname = cu.fetchone()[0]
            if seqname is None:
                # this primary key does not have a sequence associated with it
                continue
            # get the max seq value
            cu.execute("select max(%s) from %s" % (col, table))
            seqval = cu.fetchone()[0]
            if not seqval:
                seqval = 1
            else:
                seqval += 1 # we need the next one in line
            # now reset the sequence for the primary key
            cu.execute("select pg_catalog.setval(?, ?, false)", (seqname, seqval))
            ret = cu.fetchone()[0]
            assert (ret == seqval)
            if self.verbose:
                print "SETVAL %s = %d (%s.%s)" % (seqname, ret, table, col)
    # functions for when the instance is a target
    def prepareInsert(self, table, fields):
        return PgSQLLoader(self.db, table, fields)
    def finalize(self, version):
        Database.finalize(self, version)
        print "VACUUM ANALYZE"
        cu = self.db.cursor()
        cu.execute("VACUUM ANALYZE")

class MySQLDatabase(Database):
    def __init__(self, db, verbose=True):
        Database.__init__(self, "mysql", db, verbose)
    # functions for when the instance is a target
    def finalize(self, version):
        Database.finalize(self, version)
        cu = self.db.cursor()
        for t in TableList:
            print "ANALYZE", t
            cu.execute("ANALYZE LOCAL TABLE %s" %(t,))

def getdb(driver, db, verbose=True):
    if driver == "postgresql":
        return PgSQLDatabase(db, verbose)
    elif driver == "mysql":
        return MySQLDatabase(db, verbose)
    return Database(driver, db, verbose)

<commit_msg>db2db: avoid flushing the query cache when migrating from MySQL into other repos.<commit_after>#!/usr/bin/env python
#
# Copyright (c) 2007-2008 rPath, Inc.
#
# This program is distributed under the terms of the Common Public License,
# version 1.0. A copy of this license should have been distributed with this
# source file in a file called LICENSE. If it is not present, the license
# is always available at http://www.rpath.com/permanent/licenses/CPL-1.0.
#
# This program is distributed in the hope that it will be useful, but
# without any warranty; without even the implied warranty of merchantability
# or fitness for a particular purpose. See the Common Public License for
# full details.
#

# Cristian Gafton, <gafton@rpath.com>
# handles extracting and inserting into various backends for database migration

import sys

from conary import dbstore

from schema import getTables, getIndexes
from tablelist import TableList

# database loader that knows how to insert rows efficiently
class Loader:
    binaryCols = [
        "salt", "password", "data", "entitlement", "pgpkey",
        "stream", "fileid", "pathid", "sha1" ]
    def __init__(self, db, table, fields):
        self.db = db
        self.table = table
        self.sql = "INSERT INTO %s (%s) VALUES (%s)" % (
            table, ", ".join(fields), ",".join(["?"]*len(fields)))
        self.cu = self.db.cursor()
        # check the fields are all present in the target
        self.cu.execute("SELECT * FROM %s LIMIT 1" % (table,))
        set1 = set([x.lower() for x in fields])
        set2 = set([x.lower() for x in self.cu.fields()])
        assert (set1.difference(set2) == set()), "Fields differ from src to target: %s %s" %(set1, set2)
        self.fields = fields
        # make sure we know when to use cu.binary()
        self.funcs = [ self.__getfunc(x) for x in fields ]
    def announce(self):
        sys.stdout.write("1by1 load table: %s\r" % (self.table,))
        sys.stdout.flush()
    # fields which we should tag as binary
    def __getfunc(self, field):
        if field.lower() in self.binaryCols:
            return self.cu.binary
        return lambda a: a
    # fix up the values in a row to match the transforms
    def rowvals(self, row):
        return tuple(map(lambda (f, x): f(x), zip(self.funcs, row)))
    def insertRows(self, rows, callback=None):
        self.cu.executemany(self.sql, [self.rowvals(row) for row in rows])
        if callback:
            callback.increment(len(rows))
        return len(rows)

class PgSQLLoader(Loader):
    def __init__(self, db, table, fields):
        Loader.__init__(self, db, table, fields)
        self.__usebulk = False
        if hasattr(self.db.dbh, "bulkload"):
            self.__usebulk = True
        if not self.__usebulk:
            print "WARNING: not using bulk load, update your python-pgsql bindings!"
    def bulkInsert(self, rows, callback):
        self.db.dbh.bulkload(self.table, (self.rowvals(row) for row in rows), self.fields)
        if callback:
            callback.increment(len(rows))
        return len(rows)
    def insertRows(self, rows, callback=None):
        if self.__usebulk:
            return self.bulkInsert(rows, callback)
        else:
            return Loader.insertRows(self, rows, callback)

class Database:
    def __init__(self, driver, db, verbose=True):
        self.db = dbstore.connect(db, driver)
        self.driver = driver
        self.db.loadSchema()
        self.verbose = verbose
        self._hint = ''
    def createSchema(self):
        # create the tables, avoid the indexes
        cu = self.db.cursor()
        for stmt in getTables(self.driver):
            if self.verbose:
                print stmt
            cu.execute(stmt)
        self.db.loadSchema()
    def createIndexes(self):
        cu = self.db.cursor()
        for stmt in getIndexes(self.driver):
            if self.verbose:
                print stmt
            cu.execute(stmt)
        self.db.loadSchema()
    # check self.db.tables against the TableList
    def checkTablesList(self, isSrc=True):
        #  check that we are migrating all the tables in the source
        self.db.loadSchema()
        skip = ['databaseversion', 'instructionsets', 'commitlock']
        knowns = [x.lower() for x in TableList]
        haves = [x.lower() for x in self.db.tables]
        if isSrc:
            which = "Source"
        else:
            which = "Target"
        # tableList should not have items not present in the db
        onlyKnowns = set(knowns).difference(set(haves)).difference(set(skip))
        if onlyKnowns:
            raise RuntimeError("%s schema (%s) does not have table(s) %s" %(
                which, self.driver, onlyKnowns))
        # we should not have extra tables in the source
        onlyHaves = set(haves).difference(set(knowns)).difference(set(skip))
        if onlyHaves and isSrc:
            raise RuntimeError("TableList needs to be updated to handle tables", onlyHaves)
        return True

    # functions for when the instance is a source
    def getCount(self, table):
        cu = self.db.cursor()
        cu.execute("select count(*) from %s" % (table,))
        return cu.fetchall()[0][0]
    def getFields(self, table):
        cu = self.db.cursor()
        cu.execute("SELECT * FROM %s LIMIT 1" % (table,))
        return [x.lower() for x in cu.fields()]
    def getTables(self):
        return [x.lower() for x in self.db.tables]
    def iterRows(self, table, fields = "*"):
        cu = self.db.itercursor()
        cu.execute("select %s %s from %s" % (self._hint, fields, table))
        return cu
    # functions for when the instance is a target
    def prepareInsert(self, table, fields):
        return Loader(self.db, table, fields)
    def finalize(self, version):
        self.db.setVersion(version)
        self.db.commit()
    # useful shortcut
    def commit(self):
        return self.db.commit()
    def close(self):
        self.db.close()

class PgSQLDatabase(Database):
    def __init__(self, db, verbose=True):
        Database.__init__(self, "postgresql", db, verbose)
    def createIndexes(self):
        Database.createIndexes(self)
        # fix the primary keys
        self.fix_primary_keys()
        # update the primary key sequences for all tables
    def fix_primary_keys(self):
        cu = self.db.cursor()
        # get the name of the primary key
        cu.execute("""
        select
            t.relname as table_name,
            col.attname as column_name
        from pg_class t
        join pg_index i on t.oid = i.indrelid and i.indisprimary = true
        join pg_class ind on i.indexrelid = ind.oid
        join pg_attribute col on col.attrelid = t.oid and col.attnum = i.indkey[0]
        where i.indnatts = 1
          and pg_catalog.pg_table_is_visible(t.oid)
        """)
        for (t, col) in cu.fetchall():
            table = t.lower()
            cu.execute("select pg_catalog.pg_get_serial_sequence(?, ?)", (table, col))
            seqname = cu.fetchone()[0]
            if seqname is None:
                # this primary key does not have a sequence associated with it
                continue
            # get the max seq value
            cu.execute("select max(%s) from %s" % (col, table))
            seqval = cu.fetchone()[0]
            if not seqval:
                seqval = 1
            else:
                seqval += 1 # we need the next one in line
            # now reset the sequence for the primary key
            cu.execute("select pg_catalog.setval(?, ?, false)", (seqname, seqval))
            ret = cu.fetchone()[0]
            assert (ret == seqval)
            if self.verbose:
                print "SETVAL %s = %d (%s.%s)" % (seqname, ret, table, col)
    # functions for when the instance is a target
    def prepareInsert(self, table, fields):
        return PgSQLLoader(self.db, table, fields)
    def finalize(self, version):
        Database.finalize(self, version)
        cu = self.db.cursor()
        if self.verbose:
            print "VACUUM ANALYZE"
        cu.execute("VACUUM ANALYZE")

class MySQLDatabase(Database):
    def __init__(self, db, verbose=True):
        Database.__init__(self, "mysql", db, verbose)
        self._hint = "/*!40001 SQL_NO_CACHE */"
    # functions for when the instance is a target
    def finalize(self, version):
        Database.finalize(self, version)
        cu = self.db.cursor()
        if self.verbose:
            print "ANALYZE"
        for t in TableList:
            cu.execute("ANALYZE LOCAL TABLE %s" %(t,))

def getdb(driver, db, verbose=True):
    if driver == "postgresql":
        return PgSQLDatabase(db, verbose)
    elif driver == "mysql":
        return MySQLDatabase(db, verbose)
    return Database(driver, db, verbose)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
=====================
Cython related magics
=====================

Magic command interface for interactive work with Cython

.. note::

  The ``Cython`` package needs to be installed separately. It
  can be obtained using ``easy_install`` or ``pip``.

Usage
=====

To enable the magics below, execute ``%load_ext cythonmagic``.

``%%cython``

{CYTHON_DOC}

``%%cython_inline``

{CYTHON_INLINE_DOC}

``%%cython_pyximport``

{CYTHON_PYXIMPORT_DOC}

Author:
* Brian Granger

Code moved from IPython and adapted by:
* Martín Gaitán

Parts of this code were taken from Cython.inline.
"""
#-----------------------------------------------------------------------------
# Copyright (C) 2010-2011, IPython Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file ipython-COPYING.rst, distributed with this software.
#-----------------------------------------------------------------------------

from __future__ import absolute_import, print_function

import imp
import io
import os
import re
import sys
import time

try:
    reload
except NameError:   # Python 3
    from imp import reload

try:
    import hashlib
except ImportError:
    import md5 as hashlib

from distutils.core import Distribution, Extension
from distutils.command.build_ext import build_ext

from IPython.core import display
from IPython.core import magic_arguments
from IPython.core.magic import Magics, magics_class, cell_magic
from IPython.utils import py3compat
try:
    from IPython.paths import get_ipython_cache_dir
except ImportError:
    # older IPython version
    from IPython.utils.path import get_ipython_cache_dir
from IPython.utils.text import dedent

from ..Shadow import __version__ as cython_version
from ..Compiler.Errors import CompileError
from .Inline import cython_inline
from .Dependencies import cythonize

@magics_class
class CythonMagics(Magics):

    def __init__(self, shell):
        super(CythonMagics,self).__init__(shell)
        self._reloads = {}
        self._code_cache = {}
        self._pyximport_installed = False

    def _import_all(self, module):
        for k,v in module.__dict__.items():
            if not k.startswith('__'):
                self.shell.push({k:v})

    @cell_magic
    def cython_inline(self, line, cell):
        """Compile and run a Cython code cell using Cython.inline.

        This magic simply passes the body of the cell to Cython.inline
        and returns the result. If the variables `a` and `b` are defined
        in the user's namespace, here is a simple example that returns
        their sum::

            %%cython_inline
            return a+b

        For most purposes, we recommend the usage of the `%%cython` magic.
        """
        locs = self.shell.user_global_ns
        globs = self.shell.user_ns
        return cython_inline(cell, locals=locs, globals=globs)

    @cell_magic
    def cython_pyximport(self, line, cell):
        """Compile and import a Cython code cell using pyximport.

        The contents of the cell are written to a `.pyx` file in the current
        working directory, which is then imported using `pyximport`. This
        magic requires a module name to be passed::

            %%cython_pyximport modulename
            def f(x):
                return 2.0*x

        The compiled module is then imported and all of its symbols are
        injected into the user's namespace. For most purposes, we recommend
        the usage of the `%%cython` magic.
        """
        module_name = line.strip()
        if not module_name:
            raise ValueError('module name must be given')
        fname = module_name + '.pyx'
        with io.open(fname, 'w', encoding='utf-8') as f:
            f.write(cell)
        if 'pyximport' not in sys.modules or not self._pyximport_installed:
            import pyximport
            pyximport.install(reload_support=True)
            self._pyximport_installed = True
        if module_name in self._reloads:
            module = self._reloads[module_name]
            reload(module)
        else:
            __import__(module_name)
            module = sys.modules[module_name]
            self._reloads[module_name] = module
        self._import_all(module)

    @magic_arguments.magic_arguments()
    @magic_arguments.argument(
        '-c', '--compile-args', action='append', default=[],
        help="Extra flags to pass to compiler via the `extra_compile_args` "
             "Extension flag (can be specified  multiple times)."
    )
    @magic_arguments.argument(
        '--link-args', action='append', default=[],
        help="Extra flags to pass to linker via the `extra_link_args` "
             "Extension flag (can be specified  multiple times)."
    )
    @magic_arguments.argument(
        '-l', '--lib', action='append', default=[],
        help="Add a library to link the extension against (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-n', '--name',
        help="Specify a name for the Cython module."
    )
    @magic_arguments.argument(
        '-L', dest='library_dirs', metavar='dir', action='append', default=[],
        help="Add a path to the list of libary directories (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-I', '--include', action='append', default=[],
        help="Add a path to the list of include directories (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-+', '--cplus', action='store_true', default=False,
        help="Output a C++ rather than C file."
    )
    @magic_arguments.argument(
        '-f', '--force', action='store_true', default=False,
        help="Force the compilation of a new module, even if the source has been "
             "previously compiled."
    )
    @magic_arguments.argument(
        '-a', '--annotate', action='store_true', default=False,
        help="Produce a colorized HTML version of the source."
    )
    @cell_magic
    def cython(self, line, cell):
        """Compile and import everything from a Cython code cell.

        The contents of the cell are written to a `.pyx` file in the
        directory `IPYTHONDIR/cython` using a filename with the hash of the
        code. This file is then cythonized and compiled. The resulting module
        is imported and all of its symbols are injected into the user's
        namespace. The usage is similar to that of `%%cython_pyximport` but
        you don't have to pass a module name::

            %%cython
            def f(x):
                return 2.0*x

        To compile OpenMP codes, pass the required  `--compile-args`
        and `--link-args`.  For example with gcc::

            %%cython --compile-args=-fopenmp --link-args=-fopenmp
            ...
        """
        args = magic_arguments.parse_argstring(self.cython, line)
        code = cell if cell.endswith('\n') else cell+'\n'
        lib_dir = os.path.join(get_ipython_cache_dir(), 'cython')
        quiet = True
        key = code, line, sys.version_info, sys.executable, cython_version

        if not os.path.exists(lib_dir):
            os.makedirs(lib_dir)

        if args.force:
            # Force a new module name by adding the current time to the
            # key which is hashed to determine the module name.
            key += time.time(),

        if args.name:
            module_name = py3compat.unicode_to_str(args.name)
        else:
            module_name = "_cython_magic_" + hashlib.md5(str(key).encode('utf-8')).hexdigest()
        module_path = os.path.join(lib_dir, module_name + self.so_ext)

        have_module = os.path.isfile(module_path)
        need_cythonize = not have_module

        if args.annotate:
            html_file = os.path.join(lib_dir, module_name + '.html')
            if not os.path.isfile(html_file):
                need_cythonize = True

        if need_cythonize:
            c_include_dirs = args.include
            if 'numpy' in code:
                import numpy
                c_include_dirs.append(numpy.get_include())
            pyx_file = os.path.join(lib_dir, module_name + '.pyx')
            pyx_file = py3compat.cast_bytes_py2(pyx_file, encoding=sys.getfilesystemencoding())
            with io.open(pyx_file, 'w', encoding='utf-8') as f:
                f.write(code)
            extension = Extension(
                name = module_name,
                sources = [pyx_file],
                include_dirs = c_include_dirs,
                library_dirs = args.library_dirs,
                extra_compile_args = args.compile_args,
                extra_link_args = args.link_args,
                libraries = args.lib,
                language = 'c++' if args.cplus else 'c',
            )
            build_extension = self._get_build_extension()
            try:
                opts = dict(
                    quiet=quiet,
                    annotate = args.annotate,
                    force = True,
                    )
                build_extension.extensions = cythonize([extension], **opts)
            except CompileError:
                return

        if not have_module:
            build_extension.build_temp = os.path.dirname(pyx_file)
            build_extension.build_lib  = lib_dir
            build_extension.run()
            self._code_cache[key] = module_name

        module = imp.load_dynamic(module_name, module_path)
        self._import_all(module)

        if args.annotate:
            try:
                with io.open(html_file, encoding='utf-8') as f:
                    annotated_html = f.read()
            except IOError as e:
                # File could not be opened. Most likely the user has a version
                # of Cython before 0.15.1 (when `cythonize` learned the
                # `force` keyword argument) and has already compiled this
                # exact source without annotation.
                print('Cython completed successfully but the annotated '
                      'source could not be read.', file=sys.stderr)
                print(e, file=sys.stderr)
            else:
                return display.HTML(self.clean_annotated_html(annotated_html))

    @property
    def so_ext(self):
        """The extension suffix for compiled modules."""
        try:
            return self._so_ext
        except AttributeError:
            self._so_ext = self._get_build_extension().get_ext_filename('')
            return self._so_ext

    def _clear_distutils_mkpath_cache(self):
        """clear distutils mkpath cache

        prevents distutils from skipping re-creation of dirs that have been removed
        """
        try:
            from distutils.dir_util import _path_created
        except ImportError:
            pass
        else:
            _path_created.clear()

    def _get_build_extension(self):
        self._clear_distutils_mkpath_cache()
        dist = Distribution()
        config_files = dist.find_config_files()
        try:
            config_files.remove('setup.cfg')
        except ValueError:
            pass
        dist.parse_config_files(config_files)
        build_extension = build_ext(dist)
        build_extension.finalize_options()
        return build_extension

    @staticmethod
    def clean_annotated_html(html):
        """Clean up the annotated HTML source.

        Strips the link to the generated C or C++ file, which we do not
        present to the user.
        """
        r = re.compile('
Raw output: (.*)')
        html = '\n'.join(l for l in html.splitlines() if not r.match(l))
        return html

__doc__ = __doc__.format(
                # rST doesn't see the -+ flag as part of an option list, so we
                # hide it from the module-level docstring.
                CYTHON_DOC = dedent(CythonMagics.cython.__doc__\
                            .replace('-+, --cplus','--cplus    ')),
                CYTHON_INLINE_DOC = dedent(CythonMagics.cython_inline.__doc__),
                CYTHON_PYXIMPORT_DOC = dedent(CythonMagics.cython_pyximport.__doc__),
)
<commit_msg>fix typo<commit_after># -*- coding: utf-8 -*-
"""
=====================
Cython related magics
=====================

Magic command interface for interactive work with Cython

.. note::

  The ``Cython`` package needs to be installed separately. It
  can be obtained using ``easy_install`` or ``pip``.

Usage
=====

To enable the magics below, execute ``%load_ext cythonmagic``.

``%%cython``

{CYTHON_DOC}

``%%cython_inline``

{CYTHON_INLINE_DOC}

``%%cython_pyximport``

{CYTHON_PYXIMPORT_DOC}

Author:
* Brian Granger

Code moved from IPython and adapted by:
* Martín Gaitán

Parts of this code were taken from Cython.inline.
"""
#-----------------------------------------------------------------------------
# Copyright (C) 2010-2011, IPython Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file ipython-COPYING.rst, distributed with this software.
#-----------------------------------------------------------------------------

from __future__ import absolute_import, print_function

import imp
import io
import os
import re
import sys
import time

try:
    reload
except NameError:   # Python 3
    from imp import reload

try:
    import hashlib
except ImportError:
    import md5 as hashlib

from distutils.core import Distribution, Extension
from distutils.command.build_ext import build_ext

from IPython.core import display
from IPython.core import magic_arguments
from IPython.core.magic import Magics, magics_class, cell_magic
from IPython.utils import py3compat
try:
    from IPython.paths import get_ipython_cache_dir
except ImportError:
    # older IPython version
    from IPython.utils.path import get_ipython_cache_dir
from IPython.utils.text import dedent

from ..Shadow import __version__ as cython_version
from ..Compiler.Errors import CompileError
from .Inline import cython_inline
from .Dependencies import cythonize

@magics_class
class CythonMagics(Magics):

    def __init__(self, shell):
        super(CythonMagics,self).__init__(shell)
        self._reloads = {}
        self._code_cache = {}
        self._pyximport_installed = False

    def _import_all(self, module):
        for k,v in module.__dict__.items():
            if not k.startswith('__'):
                self.shell.push({k:v})

    @cell_magic
    def cython_inline(self, line, cell):
        """Compile and run a Cython code cell using Cython.inline.

        This magic simply passes the body of the cell to Cython.inline
        and returns the result. If the variables `a` and `b` are defined
        in the user's namespace, here is a simple example that returns
        their sum::

            %%cython_inline
            return a+b

        For most purposes, we recommend the usage of the `%%cython` magic.
        """
        locs = self.shell.user_global_ns
        globs = self.shell.user_ns
        return cython_inline(cell, locals=locs, globals=globs)

    @cell_magic
    def cython_pyximport(self, line, cell):
        """Compile and import a Cython code cell using pyximport.

        The contents of the cell are written to a `.pyx` file in the current
        working directory, which is then imported using `pyximport`. This
        magic requires a module name to be passed::

            %%cython_pyximport modulename
            def f(x):
                return 2.0*x

        The compiled module is then imported and all of its symbols are
        injected into the user's namespace. For most purposes, we recommend
        the usage of the `%%cython` magic.
        """
        module_name = line.strip()
        if not module_name:
            raise ValueError('module name must be given')
        fname = module_name + '.pyx'
        with io.open(fname, 'w', encoding='utf-8') as f:
            f.write(cell)
        if 'pyximport' not in sys.modules or not self._pyximport_installed:
            import pyximport
            pyximport.install(reload_support=True)
            self._pyximport_installed = True
        if module_name in self._reloads:
            module = self._reloads[module_name]
            reload(module)
        else:
            __import__(module_name)
            module = sys.modules[module_name]
            self._reloads[module_name] = module
        self._import_all(module)

    @magic_arguments.magic_arguments()
    @magic_arguments.argument(
        '-c', '--compile-args', action='append', default=[],
        help="Extra flags to pass to compiler via the `extra_compile_args` "
             "Extension flag (can be specified  multiple times)."
    )
    @magic_arguments.argument(
        '--link-args', action='append', default=[],
        help="Extra flags to pass to linker via the `extra_link_args` "
             "Extension flag (can be specified  multiple times)."
    )
    @magic_arguments.argument(
        '-l', '--lib', action='append', default=[],
        help="Add a library to link the extension against (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-n', '--name',
        help="Specify a name for the Cython module."
    )
    @magic_arguments.argument(
        '-L', dest='library_dirs', metavar='dir', action='append', default=[],
        help="Add a path to the list of library directories (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-I', '--include', action='append', default=[],
        help="Add a path to the list of include directories (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-+', '--cplus', action='store_true', default=False,
        help="Output a C++ rather than C file."
    )
    @magic_arguments.argument(
        '-f', '--force', action='store_true', default=False,
        help="Force the compilation of a new module, even if the source has been "
             "previously compiled."
    )
    @magic_arguments.argument(
        '-a', '--annotate', action='store_true', default=False,
        help="Produce a colorized HTML version of the source."
    )
    @cell_magic
    def cython(self, line, cell):
        """Compile and import everything from a Cython code cell.

        The contents of the cell are written to a `.pyx` file in the
        directory `IPYTHONDIR/cython` using a filename with the hash of the
        code. This file is then cythonized and compiled. The resulting module
        is imported and all of its symbols are injected into the user's
        namespace. The usage is similar to that of `%%cython_pyximport` but
        you don't have to pass a module name::

            %%cython
            def f(x):
                return 2.0*x

        To compile OpenMP codes, pass the required  `--compile-args`
        and `--link-args`.  For example with gcc::

            %%cython --compile-args=-fopenmp --link-args=-fopenmp
            ...
        """
        args = magic_arguments.parse_argstring(self.cython, line)
        code = cell if cell.endswith('\n') else cell+'\n'
        lib_dir = os.path.join(get_ipython_cache_dir(), 'cython')
        quiet = True
        key = code, line, sys.version_info, sys.executable, cython_version

        if not os.path.exists(lib_dir):
            os.makedirs(lib_dir)

        if args.force:
            # Force a new module name by adding the current time to the
            # key which is hashed to determine the module name.
            key += time.time(),

        if args.name:
            module_name = py3compat.unicode_to_str(args.name)
        else:
            module_name = "_cython_magic_" + hashlib.md5(str(key).encode('utf-8')).hexdigest()
        module_path = os.path.join(lib_dir, module_name + self.so_ext)

        have_module = os.path.isfile(module_path)
        need_cythonize = not have_module

        if args.annotate:
            html_file = os.path.join(lib_dir, module_name + '.html')
            if not os.path.isfile(html_file):
                need_cythonize = True

        if need_cythonize:
            c_include_dirs = args.include
            if 'numpy' in code:
                import numpy
                c_include_dirs.append(numpy.get_include())
            pyx_file = os.path.join(lib_dir, module_name + '.pyx')
            pyx_file = py3compat.cast_bytes_py2(pyx_file, encoding=sys.getfilesystemencoding())
            with io.open(pyx_file, 'w', encoding='utf-8') as f:
                f.write(code)
            extension = Extension(
                name = module_name,
                sources = [pyx_file],
                include_dirs = c_include_dirs,
                library_dirs = args.library_dirs,
                extra_compile_args = args.compile_args,
                extra_link_args = args.link_args,
                libraries = args.lib,
                language = 'c++' if args.cplus else 'c',
            )
            build_extension = self._get_build_extension()
            try:
                opts = dict(
                    quiet=quiet,
                    annotate = args.annotate,
                    force = True,
                    )
                build_extension.extensions = cythonize([extension], **opts)
            except CompileError:
                return

        if not have_module:
            build_extension.build_temp = os.path.dirname(pyx_file)
            build_extension.build_lib  = lib_dir
            build_extension.run()
            self._code_cache[key] = module_name

        module = imp.load_dynamic(module_name, module_path)
        self._import_all(module)

        if args.annotate:
            try:
                with io.open(html_file, encoding='utf-8') as f:
                    annotated_html = f.read()
            except IOError as e:
                # File could not be opened. Most likely the user has a version
                # of Cython before 0.15.1 (when `cythonize` learned the
                # `force` keyword argument) and has already compiled this
                # exact source without annotation.
                print('Cython completed successfully but the annotated '
                      'source could not be read.', file=sys.stderr)
                print(e, file=sys.stderr)
            else:
                return display.HTML(self.clean_annotated_html(annotated_html))

    @property
    def so_ext(self):
        """The extension suffix for compiled modules."""
        try:
            return self._so_ext
        except AttributeError:
            self._so_ext = self._get_build_extension().get_ext_filename('')
            return self._so_ext

    def _clear_distutils_mkpath_cache(self):
        """clear distutils mkpath cache

        prevents distutils from skipping re-creation of dirs that have been removed
        """
        try:
            from distutils.dir_util import _path_created
        except ImportError:
            pass
        else:
            _path_created.clear()

    def _get_build_extension(self):
        self._clear_distutils_mkpath_cache()
        dist = Distribution()
        config_files = dist.find_config_files()
        try:
            config_files.remove('setup.cfg')
        except ValueError:
            pass
        dist.parse_config_files(config_files)
        build_extension = build_ext(dist)
        build_extension.finalize_options()
        return build_extension

    @staticmethod
    def clean_annotated_html(html):
        """Clean up the annotated HTML source.

        Strips the link to the generated C or C++ file, which we do not
        present to the user.
        """
        r = re.compile('
Raw output: (.*)')
        html = '\n'.join(l for l in html.splitlines() if not r.match(l))
        return html

__doc__ = __doc__.format(
                # rST doesn't see the -+ flag as part of an option list, so we
                # hide it from the module-level docstring.
                CYTHON_DOC = dedent(CythonMagics.cython.__doc__\
                            .replace('-+, --cplus','--cplus    ')),
                CYTHON_INLINE_DOC = dedent(CythonMagics.cython_inline.__doc__),
                CYTHON_PYXIMPORT_DOC = dedent(CythonMagics.cython_pyximport.__doc__),
)
<|endoftext|>