<commit_before>from dateutil.parser import parse as _parse

def _nbaLeague(x):
    leagues = {"NBA":"00", "WNBA":"10", "NBADL":"20"}
    try:
        return leagues[x.upper()]
    except:
        raise Exception("Please use one of the following values for League: 'NBA', 'WNBA', 'NBADL'")
def _nbaSeason(x):
    if str(x) == "1999":
        return "1999-00"
    elif len(str(x)) == 4:
        try:
            return "-".join([str(x),str(int(x) % 100 + 1)])
        except ValueError:
            raise Exception("Enter the four digit year for the first half of the desired season")
    else: raise Exception("Enter the four digit year for the first half of the desired season")
def _seasonID(x):
    if len(str(x)) == 4:
        try:
            return "".join(["2",str(x)])
        except ValueError:
            raise Exception("Enter the four digit year for the first half of the desired season")
    else: raise Exception("Enter the four digit year for the first half of the desired season")
def _measureType(x):
    measure = {1:"Base", 2:"Advanced", 3:"Misc", 4:"Four Factors", 5:"Scoring", 6:"Opponent", 7:"Usage"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _Scope(x):
    measure = {1:'',2:"RS",3:"S",4:"Rookies"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PerModeSmall48(x):
    measure = {1:"Totals",2:"PerGame",3:"Per48"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PerModeSmall36(x):
    measure = {1:"Totals",2:"PerGame",3:"Per36"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PerModeMini(x):
    measure = {1:"Totals",2:"PerGame"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))

def _AheadBehind(x):
    measure = {1:"Ahead or Behind",2:"Behind or Tied",3:"Ahead or Tied"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _ClutchTime(x):
    measure = {1:"Last 5 Minutes",2:"Last 4 Minutes",3:"Last 3 Minutes",
               3:"Last 2 Minutes",4:"Last 1 Minute",5:"Last 30 Seconds",
               6:"Last 10 Seconds", 7:''}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _GameScope(x):
    measure = {1:'', 2:"Yesterday", 3:"Last 10"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PlayerExperience(x):
    measure = {1:'', 2:"Rookie",3:"Sophomore",4:"Veteran"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PlayerPosition(x):
    measure = {1:'',2:"F", 3:"C", 4:"G", 5:"C-F", 6:"F-C", 7:"F-G", 8:"G-F"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _StarterBench(x):
    measure = {1:'', 2:"Starters", 3:"Bench"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PerModeLarge(x):
    measure = {1:"Totals",2:"PerGame",3:"MinutesPer",4:"Per48",5:"Per40",
               6:"Per36",7:"PerMinute", 8:"PerPossession", 9:"PerPlay",
               10:"Per100Possessions", 11:"Per100Plays"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PlusMinus(x):
    measure = {2:"Y",1:"N"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PaceAdjust(x):
    measure = {2:"Y",1:"N"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _Rank(x):
    measure = {2:"Y",1:"N"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _SeasonType(x):
    measure = {1:"Regular Season",2:"Playoffs",3:"All Star"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _Outcome(x):
    measure = {1:'',2:"W",3:"L"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _Location(x):
    measure = {1:'', 2:"Home",3:"Road"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _SeasonSegment(x):
    measure = {1:'', 2:"Post All-Star",3:"Pre All-Star"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _VsConference(x):
    measure = {1:'',2:"East",3:"West"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _VsDivision(x):
    measure = {1:'',2:"Atlantic",3:"Central",4:"Northwest",5:"Pacific",6:"Southeast",7:"Southwest",8:"East",9:"West"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _GameSegment(x):
    measure = {1:'',2:"First Half",3:"Overtime",2:"Second Half"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _DistanceRange(x):
    measure = {1:"5ft Range",2:"8ft Range",3:"By Zone"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _valiDate(date_text):
    try:
        _parse(date_text)
    except ValueError:
        raise ValueError("Incorrect data format, should be YYYY-MM-DD or MM-DD-YYYY")
def _ContextMeasure(x):
    measure = { 1:'FGM',2:'FGA',3:"FG_PCT",4:'FG3M',5:'FG3A',6:'FG3_PCT',8:'PF',
                9:'EFG_PCT',10:'TS_PCT',11:'PTS_FB',12:'PTS_OFF_TOV',
                13:'PTS_2ND_CHANCE',14:'PF'}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _Position(x):
    measure = {1:'',2:'Guard',3:'Center',4:'Forward'}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
# BoxScore -- RangeType must be between 0 and 2.
__all__ = ['_Position', '_ContextMeasure', '_valiDate', '_DistanceRange',
           '_GameSegment', '_VsDivision', '_VsConference', '_SeasonSegment'
           '_Location', '_Outcome', '_SeasonType', '_Rank', '_PaceAdjust',
           '_PlusMinus', '_PerModeSmall48', '_StarterBench', '_PlayerPosition',
           '_PlayerExperience', '_GameScope', '_ClutchTime', '_AheadBehind',
           '_PerModeLarge', '_Scope', '_measureType', '_seasonID', '_nbaSeason',
           '_nbaLeague','_PerModeSmall36', '_PerModeMini']<commit_msg>bug fixes, added handling for blank values in _valiDate<commit_after>from dateutil.parser import parse as _parse

def _nbaLeague(x):
    """Takes in initials of league and returns numeric API Code

    Input Values: "NBA", "WNBA", or "NBADL"

    Used in: _Draft.Anthro(), _Draft.Agility(), _Draft.NonStationaryShooting(),
    _Draft.SpotUpShooting(), _Draft.Combine()

    """
    leagues = {"NBA":"00", "WNBA":"10", "NBADL":"20"}
    try:
        return leagues[x.upper()]
    except:
        raise Exception("Please use one of the following values for League: 'NBA', 'WNBA', 'NBADL'")
def _nbaSeason(x):
    """Takes in 4-digit year for first half of season and returns API appropriate formatted code

    Input Values: YYYY

    Used in: _Draft.Anthro(), _Draft.Agility(), _Draft.NonStationaryShooting(),
    _Draft.SpotUpShooting(), _Draft.Combine()

    """
    if str(x) == "1999":
        return "1999-00"
    elif len(str(x)) == 4:
        try:
            return "-".join([str(x),str(int(x) % 100 + 1)])
        except ValueError:
            raise Exception("Enter the four digit year for the first half of the desired season")
    else: raise Exception("Enter the four digit year for the first half of the desired season")
def _seasonID(x):
    """takes in 4-digit years and returns API formatted seasonID

    Input Values: YYYY

    Used in:

    """
    if len(str(x)) == 4:
        try:
            return "".join(["2",str(x)])
        except ValueError:
            raise Exception("Enter the four digit year for the first half of the desired season")
    else: raise Exception("Enter the four digit year for the first half of the desired season")
def _measureType(x):
    """Takes Numeric Code and returns String API code

    Input Values: 1:"Base", 2:"Advanced", 3:"Misc", 4:"Four Factors", 5:"Scoring", 6:"Opponent", 7:"Usage"

    Used in:

    """
    measure = {1:"Base", 2:"Advanced", 3:"Misc", 4:"Four Factors", 5:"Scoring", 6:"Opponent", 7:"Usage"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _Scope(x):
    """Takes Numeric Code and returns String API code

    Input Values: 1:"", 2:"RS", 3:"S", 4:"Rookies"
    Used in:

    """
    measure = {1:'',2:"RS",3:"S",4:"Rookies"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PerModeSmall48(x):
    """Takes Numeric Code and returns String API code

    Input Values: 1:"Base", 2:"Advanced", 3:"Misc", 4:"Four Factors", 5:"Scoring", 6:"Opponent", 7:"Usage"

    Used in:

    """
    measure = {1:"Totals",2:"PerGame",3:"Per48"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PerModeSmall36(x):
    """Takes Numeric Code and returns String API code

    Input Values: 1:"Base", 2:"Advanced", 3:"Misc", 4:"Four Factors", 5:"Scoring", 6:"Opponent", 7:"Usage"

    Used in:

    """
    measure = {1:"Totals",2:"PerGame",3:"Per36"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PerModeMini(x):
    """Takes Numeric Code and returns String API code

    Input Values: 1:"Totals", 2:"PerGame"

    Used in:

    """
    measure = {1:"Totals",2:"PerGame"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PerModeLarge(x):
    measure = {1:"Totals",2:"PerGame",3:"MinutesPer",4:"Per48",5:"Per40",
               6:"Per36",7:"PerMinute", 8:"PerPossession", 9:"PerPlay",
               10:"Per100Possessions", 11:"Per100Plays"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _AheadBehind(x):
    measure = {1:"Ahead or Behind",2:"Behind or Tied",3:"Ahead or Tied"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _ClutchTime(x):
    measure = {1:"Last 5 Minutes",2:"Last 4 Minutes",3:"Last 3 Minutes",
               4:"Last 2 Minutes",5:"Last 1 Minute",6:"Last 30 Seconds",
               7:"Last 10 Seconds", 8:''}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _GameScope(x):
    measure = {1:'', 2:"Yesterday", 3:"Last 10"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PlayerExperience(x):
    measure = {1:'', 2:"Rookie",3:"Sophomore",4:"Veteran"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PlayerPosition(x):
    measure = {1:'',2:"F", 3:"C", 4:"G", 5:"C-F", 6:"F-C", 7:"F-G", 8:"G-F"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _StarterBench(x):
    measure = {1:'', 2:"Starters", 3:"Bench"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PlusMinus(x):
    measure = {2:"Y",1:"N"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _PaceAdjust(x):
    measure = {2:"Y",1:"N"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _Rank(x):
    measure = {2:"Y",1:"N"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _SeasonType(x):
    measure = {1:"Regular Season",2:"Playoffs",3:"All Star"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _SeasonType4(x):
    measure = {1:"Regular Season",2:"Playoffs",3:"All Star",4:"Pre Season"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _Outcome(x):
    measure = {1:'',2:"W",3:"L"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _Location(x):
    measure = {1:'', 2:"Home",3:"Road"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _SeasonSegment(x):
    measure = {1:'', 2:"Post All-Star",3:"Pre All-Star"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _VsConference(x):
    measure = {1:'',2:"East",3:"West"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _VsDivision(x):
    measure = {1:'',2:"Atlantic",3:"Central",4:"Northwest",5:"Pacific",6:"Southeast",7:"Southwest",8:"East",9:"West"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _GameSegment(x):
    measure = {1:'',2:"First Half",3:"Overtime",2:"Second Half"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _DistanceRange(x):
    measure = {1:"5ft Range",2:"8ft Range",3:"By Zone"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _valiDate(date_text):
    if not dateto == '':
        try:
            date = _parse(date_text)
            return str(date.date())
        except ValueError:
            raise ValueError("Incorrect data format, should be YYYY-MM-DD or MM-DD-YYYY")
    else: pass
def _ContextMeasure(x):
    measure = { 1:'FGM',2:'FGA',3:"FG_PCT",4:'FG3M',5:'FG3A',6:'FG3_PCT',8:'PF',
                9:'EFG_PCT',10:'TS_PCT',11:'PTS_FB',12:'PTS_OFF_TOV',
                13:'PTS_2ND_CHANCE',14:'PF'}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _Position(x):
    measure = {1:'',2:'Guard',3:'Center',4:'Forward'}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
def _StatCategory(x):
    measure = {1:"MIN",2:"FGM",3:"FGA",4:"FG_PCT",5:"FG3M",6:"FG3A",7:"FG3_PCT",
    8:"FTM",9:"FTA",10:"FT_PCT",11:"OREB",12:"DREB",13:"REB",14:"AST",15:"STL",
    16:"BLK",17:"TOV",18:"PTS",19:"EFF"}
    try:
        return measure[x]
    except:
        raise Exception("Please enter a number between 1 and "+str(len(measure)))
# BoxScore -- RangeType must be between 0 and 2.
__all__ = ['_Position', '_ContextMeasure', '_valiDate', '_DistanceRange',
           '_GameSegment', '_VsDivision', '_VsConference', '_SeasonSegment'
           '_Location', '_Outcome', '_SeasonType', '_Rank', '_PaceAdjust',
           '_PlusMinus', '_PerModeSmall48', '_StarterBench', '_PlayerPosition',
           '_PlayerExperience', '_GameScope', '_ClutchTime', '_AheadBehind',
           '_PerModeLarge', '_Scope', '_measureType', '_seasonID', '_nbaSeason',
           '_nbaLeague','_PerModeSmall36', '_PerModeMini']<|endoftext|>


<commit_before>"""Exports and imports mesh data"""

import bpy
import nose.tools
import math
import mathutils

from test import Base, SingleNif
from pyffi.formats.nif import NifFormat

class TestBaseGeometry(SingleNif):
    n_name = "geometry/base_geometry"
    b_cube = "Cube"

    def rotation_matrix(self):
        b_trans_mat = mathutils.Matrix.Translation((20.0, 20.0, 20.0))
        b_rot_mat = mathutils.Matrix.Rotation(math.radians(45.0), 4, 'X')
        b_scale_mat = mathutils.Matrix.Scale(0.75, 4)
        b_transform_mat = b_trans_mat * b_rot_mat * b_scale_mat
        return b_transform_mat

    def b_create_objects(self):

        #grab the last added object, avoids name confusion
        bpy.ops.mesh.primitive_cube_add()
        b_obj = bpy.data.objects[bpy.context.active_object.name]
        b_obj.name = self.b_cube
        SingleNif.b_name.append(self.b_cube) #add to new object list

        #add transformation
        b_obj.matrix_local = self.rotation_matrix()

        # primitive_cube_add sets double sided; fix this
        b_obj.data.show_double_sided = False

        #bpy.ops.wm.save_mainfile(filepath="test/autoblend/" + self.n_name)

    def b_check_data(self):
        b_obj = bpy.data.objects[self.b_cube]

        #transfroms
        b_loc, b_rot, b_scale = b_obj.matrix_local.decompose()
        #compare location value
        nose.tools.assert_equal(b_obj.location, mathutils.Vector((20.0, 20.0, 20.0)))
        #compare rotations, float comparison
        nose.tools.assert_equal(
                (b_rot.to_axis_angle()[1] - math.radians(45.0))
                < self.EPSILON, True)
        #compare scale: float comparison
        nose.tools.assert_equal(
                (b_obj.scale - mathutils.Vector((0.75, 0.75, 0.75)))
                < mathutils.Vector((self.EPSILON,self.EPSILON,self.EPSILON)), True)

        b_mesh = b_obj.data
        self.b_check_geom(b_mesh)

    def b_check_geom(self, b_mesh):
        num_triangles = len(
            [face for face in b_mesh.faces if len(face.vertices) == 3]) #check for tri
        num_triangles += 2 * len(
            [face for face in b_mesh.faces if len(face.vertices) == 4]) #face = 2 tris
        nose.tools.assert_equal(len(b_mesh.vertices), 8)
        nose.tools.assert_equal(num_triangles, 12)

    def n_check_data(self, n_data):
        n_geom = n_data.roots[0].children[0]
        nose.tools.assert_is_instance(n_geom, NifFormat.NiTriShape)
        self.n_check_geom_data(n_geom)

    def n_check_geom_data(self, n_geom):
        nose.tools.assert_equal(n_geom.data.num_vertices, 8)
        nose.tools.assert_equal(n_geom.data.num_triangles, 12)

    '''
        TODO: Additional checks needed.
        TriShape
            transforms

        TriData
            Flags: blender exports| Continue, Maya| Triangles, Pyffi| Bound.
            Consistancy:
            radius:
    '''

class TestBaseUV(TestBaseGeometry):
    n_name = "geometry/base_uv"
    b_cube = 'Cube'

    def b_create_object(self):
        TestBaseGeometry.b_create_objects()
        b_obj = bpy.data.objects[self.b_name]

        #project UV
        bpy.ops.object.mode_set(mode='EDIT', toggle=False) #ensure we are in the mode.
        bpy.ops.uv.cube_project() # named 'UVTex'
        bpy.ops.object.mode_set(mode='OBJECT', toggle=False)

        #bpy.ops.wm.save_mainfile(filepath="test/autoblend/" + self.n_name)

    def b_check_data(self):
        TestBaseGeometry.b_check_data(self)
        pass

        b_obj = bpy.data.objects[self.b_cube]
        b_mesh = b_obj.data
        nose.tools.assert_equal(len(b_mesh.uv_textures), 1)
        nose.tools.assert_equal()

        #TODO_3.0 - Separate out the UV writing from requiring a texture.

    def n_check_data(self, n_data):
        TestBaseGeometry.n_check_data(self, n_data)
        pass

        #TODO_3.0 - See above
        n_geom = n_data.roots[0].children[0]
        nose.tools.assert_equal(len(n_geom.data.uv_sets), 1)
        nose.tools.assert_equal(len(n_geom.data.uv_sets[0]), len(n_geom.data.vertices))

class TestNonUniformlyScaled(Base):
    def setup(self):
        # create a non-uniformly scaled cube
        bpy.ops.mesh.primitive_cube_add()
        b_obj = bpy.data.objects["Cube"]
        b_obj.scale = (1, 2, 3)

    @nose.tools.raises(Exception)
    def test_export(self):
        bpy.ops.export_scene.nif(
            filepath="test/export/non_uniformly_scaled_cube.nif",
            log_level='DEBUG',
            )

<commit_msg>Started nif specific tests for transformas.<commit_after>"""Exports and imports mesh data"""

import bpy
import nose.tools
import math
import mathutils

from test import Base, SingleNif
from pyffi.formats.nif import NifFormat

class TestBaseGeometry(SingleNif):
    n_name = "geometry/base_geometry"
    b_cube = "Cube"

    def rotation_matrix(self):
        b_trans_mat = mathutils.Matrix.Translation((20.0, 20.0, 20.0))
        b_rot_mat = mathutils.Matrix.Rotation(math.radians(45.0), 4, 'X')
        b_scale_mat = mathutils.Matrix.Scale(0.75, 4)
        b_transform_mat = b_trans_mat * b_rot_mat * b_scale_mat
        return b_transform_mat

    def b_create_objects(self):

        #grab the last added object, avoids name confusion
        bpy.ops.mesh.primitive_cube_add()
        b_obj = bpy.data.objects[bpy.context.active_object.name]
        b_obj.name = self.b_cube
        SingleNif.b_name.append(self.b_cube) #add to new object list

        #add transformation
        b_obj.matrix_local = self.rotation_matrix()

        # primitive_cube_add sets double sided; fix this
        b_obj.data.show_double_sided = False

        #bpy.ops.wm.save_mainfile(filepath="test/autoblend/" + self.n_name)

    def b_check_data(self):
        b_obj = bpy.data.objects[self.b_cube]

        #check transforms
        b_loc, b_rot, b_scale = b_obj.matrix_local.decompose()

        nose.tools.assert_equal(b_obj.location, mathutils.Vector((20.0, 20.0, 20.0))) #location
        nose.tools.assert_equal(
                (b_rot.to_axis_angle()[1] - math.radians(45.0))
                < self.EPSILON, True)#rotations
        nose.tools.assert_equal(
                (b_obj.scale - mathutils.Vector((0.75, 0.75, 0.75)))
                < mathutils.Vector((self.EPSILON,self.EPSILON,self.EPSILON)), True) #scale

        b_mesh = b_obj.data
        self.b_check_geom(b_mesh)

    def b_check_geom(self, b_mesh):
        num_triangles = len(
            [face for face in b_mesh.faces if len(face.vertices) == 3]) #check for tri
        num_triangles += 2 * len(
            [face for face in b_mesh.faces if len(face.vertices) == 4]) #face = 2 tris
        nose.tools.assert_equal(len(b_mesh.vertices), 8)
        nose.tools.assert_equal(num_triangles, 12)

    def n_check_data(self, n_data):
        n_geom = n_data.roots[0].children[0]
        nose.tools.assert_is_instance(n_geom, NifFormat.NiTriShape)

        #check transforms
        nose.tools.assert_equal(n_geom.translation.as_tuple(),(20.0, 20.0, 20.0))
        nose.tools.assert_equal(
                mathutils.Matrix(n_geom.rotation.as_tuple()).to_4x4(),
                mathutils.Matrix.Rotation(math.radians(45), 4, 'X'))
        print(n_geom.scale)
        nose.tools.assert_equal(
                n_geom.scale, 0.75)

        self.n_check_geom_data(n_geom)

    def n_check_geom_data(self, n_geom):
        nose.tools.assert_equal(n_geom.data.num_vertices, 8)
        nose.tools.assert_equal(n_geom.data.num_triangles, 12)

    '''
        TODO: Additional checks needed.
        TriShape
            transforms

        TriData
            Flags: blender exports| Continue, Maya| Triangles, Pyffi| Bound.
            Consistancy:
            radius:
    '''
'''
class TestBaseUV(TestBaseGeometry):
    n_name = "geometry/base_uv"
    b_cube = 'Cube'

    def b_create_object(self):
        TestBaseGeometry.b_create_objects()
        b_obj = bpy.data.objects[self.b_name]

        #project UV
        bpy.ops.object.mode_set(mode='EDIT', toggle=False) #ensure we are in the mode.
        bpy.ops.uv.cube_project() # named 'UVTex'
        bpy.ops.object.mode_set(mode='OBJECT', toggle=False)

        #bpy.ops.wm.save_mainfile(filepath="test/autoblend/" + self.n_name)

    def b_check_data(self):
        TestBaseGeometry.b_check_data(self)
        pass

        b_obj = bpy.data.objects[self.b_cube]
        b_mesh = b_obj.data
        nose.tools.assert_equal(len(b_mesh.uv_textures), 1)
        nose.tools.assert_equal()

        #TODO_3.0 - Separate out the UV writing from requiring a texture.

    def n_check_data(self, n_data):
        TestBaseGeometry.n_check_data(self, n_data)
        pass

        #TODO_3.0 - See above
        n_geom = n_data.roots[0].children[0]
        nose.tools.assert_equal(len(n_geom.data.uv_sets), 1)
        nose.tools.assert_equal(len(n_geom.data.uv_sets[0]), len(n_geom.data.vertices))
'''

class TestNonUniformlyScaled(Base):
    def setup(self):
        # create a non-uniformly scaled cube
        bpy.ops.mesh.primitive_cube_add()
        b_obj = bpy.data.objects["Cube"]
        b_obj.scale = (1, 2, 3)

    @nose.tools.raises(Exception)
    def test_export(self):
        bpy.ops.export_scene.nif(
            filepath="test/export/non_uniformly_scaled_cube.nif",
            log_level='DEBUG',
            )

<|endoftext|>


<commit_before>"""
Main URL settings page. See https://docs.djangoproject.com/en/1.10/topics/http/urls/ for more information.
"""

from django.conf import settings
from django.conf.urls import url, include
from django.conf.urls.static import static
from django.contrib import admin
from django.contrib.auth.decorators import login_required
from django.contrib.sitemaps import views
from django.http import HttpResponse
from graphene_django.views import GraphQLView

from .schema import *
from .sitemaps import Sitemap, xsl_content_type
from .utils import *

sitemaps = {
    'pages': Sitemap(
        ['index'],
        [1.0],
        ['monthly'],
        [datetime.date.today()]
    ),

    'other': Sitemap(
        ['repo', 'change-log'],
        [0.5, 1.0],
        ['monthly', 'monthly'],
        [custom_date('10/01/2017'), github_date_time_format(get_version()['published_at'])]
    )
}

admin.autodiscover()
admin.site.login = login_required(admin.site.login)

urlpatterns = [
                  url(r'^', include('viewer.urls')),
                  url(r'^', include('editor.urls'), name='editor_urls'),
                  url(r'^admin/', admin.site.urls, name='admin_urls'),
                  url(r'^', include('authy_me.urls'), name="login"),
                  url(r'^robots.txt',
                      lambda x: HttpResponse(
                          "Sitemap: https://www.gollahalli.com/sitemap.xml\nUser-agent: *\nDisallow: "
                          "/admin/\nDisallow: /cdn-cgi/\nDisallow: /portal/",
                          content_type="text/plain"), name="robots_file"),
                  url(r'^sitemap\.xml$', views.index, {'sitemaps': sitemaps, 'template_name': 'sitemap-index.xml'}),
                  url(r'^sitemap-(?P<section>.+).xml$', views.sitemap,
                      {'sitemaps': sitemaps, 'template_name': 'sitemap.xml'},
                      name='django.contrib.sitemaps.views.sitemap'),
                  url(r'^graphql', GraphQLView.as_view(graphiql=True, schema=query)),
                  url(r'^sitemap\.xsl', xsl_content_type, name='sitemap_xsl')
              ] + static(settings.MEDIA_URL, document_root=settings.MEDIA_ROOT)

handler404 = 'viewer.views.page_not_found'
handler500 = 'viewer.views.server_error'
<commit_msg>relative imports removed<commit_after>"""
Main URL settings page. See https://docs.djangoproject.com/en/1.10/topics/http/urls/ for more information.
"""

from django.conf import settings
from django.conf.urls import url, include
from django.conf.urls.static import static
from django.contrib import admin
from django.contrib.auth.decorators import login_required
from django.contrib.sitemaps import views
from django.http import HttpResponse
from graphene_django.views import GraphQLView

from gollahalli_com.schema import query
from gollahalli_com.sitemaps import Sitemap, xsl_content_type
from gollahalli_com.utils import *

sitemaps = {
    'pages': Sitemap(
        ['index'],
        [1.0],
        ['monthly'],
        [datetime.date.today()]
    ),

    'other': Sitemap(
        ['repo', 'change-log'],
        [0.5, 1.0],
        ['monthly', 'monthly'],
        [custom_date('10/01/2017'), github_date_time_format(get_version()['published_at'])]
    )
}

admin.autodiscover()
admin.site.login = login_required(admin.site.login)

urlpatterns = [
                  url(r'^', include('viewer.urls')),
                  url(r'^', include('editor.urls'), name='editor_urls'),
                  url(r'^admin/', admin.site.urls, name='admin_urls'),
                  url(r'^', include('authy_me.urls'), name="login"),
                  url(r'^robots.txt',
                      lambda x: HttpResponse(
                          "Sitemap: https://www.gollahalli.com/sitemap.xml\nUser-agent: *\nDisallow: "
                          "/admin/\nDisallow: /cdn-cgi/\nDisallow: /portal/",
                          content_type="text/plain"), name="robots_file"),
                  url(r'^sitemap\.xml$', views.index, {'sitemaps': sitemaps, 'template_name': 'sitemap-index.xml'}),
                  url(r'^sitemap-(?P<section>.+).xml$', views.sitemap,
                      {'sitemaps': sitemaps, 'template_name': 'sitemap.xml'},
                      name='django.contrib.sitemaps.views.sitemap'),
                  url(r'^graphql', GraphQLView.as_view(graphiql=True, schema=query)),
                  url(r'^sitemap\.xsl', xsl_content_type, name='sitemap_xsl')
              ] + static(settings.MEDIA_URL, document_root=settings.MEDIA_ROOT)

handler404 = 'viewer.views.page_not_found'
handler500 = 'viewer.views.server_error'
<|endoftext|>


<commit_before># It is important to store the version number in a separate file
# so that we can read it from setup.py without importing the package
__version__ = "3.0b3"
<commit_msg>Bump version to 3.0b4<commit_after># It is important to store the version number in a separate file
# so that we can read it from setup.py without importing the package
__version__ = "3.0b4"
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
This module implements the classes that deal with math.

..  :copyright: (c) 2014 by Jelte Fennema.
    :license: MIT, see License for more details.
"""

from .base_classes import LatexObject, Command, Container
from pylatex.package import Package

class Math(Container):

    """A class representing a math environment."""

    def __init__(self, data=None, inline=False):
        """.

        Args
        ----
        data: list
            Content of the math container.
        inline: bool
            If the math should be displayed inline or not.
        """

        self.inline = inline
        super().__init__(data)

    def dumps(self):
        """Return a LaTeX formatted string representing the object.

        Returns
        -------
        str

        """

        if self.inline:
            string = '$' + super().dumps(token=' ') + '$'
        else:
            string = '$$' + super().dumps(token=' ') + '$$\n'

        super().dumps()

        return string

class VectorName(Command):

    """A class representing a named vector."""

    def __init__(self, name):
        """.

        Args
        ----
        name: str
            Name of the vector
        """

        super().__init__('mathbf', arguments=name)

class Matrix(LatexObject):

    """A class representing a matrix."""

    # TODO: Convert this to an environment

    def __init__(self, matrix, mtype='p', alignment=None):
        r""".

        Args
        ----
        matrix: `numpy.ndarray` instance
            The matrix to display
        mtype: str
            What kind of brackets are used around the matrix. The different
            options and their corresponding brackets are:
            p = ( ), b = [ ], B = { }, v = \| \|, V = \|\| \|\|
        alignment: str
            How to align the content of the cells in the matrix. This is ``c``
            by default.

        References
        ----------
        * https://en.wikibooks.org/wiki/LaTeX/Mathematics#Matrices_and_arrays
        """

        import numpy as np
        self._np = np

        self.mtype = mtype
        self.matrix = matrix
        self.alignment = alignment

        super().__init__(packages=[Package('amsmath')])

    def dumps(self):
        """Return a string representin the matrix in LaTeX syntax.

        Returns
        -------
        str
        """

        string = r'\begin{'
        mtype = self.mtype + 'matrix'

        if self.alignment is not None:
            mtype += '*'
            alignment = '{' + self.alignment + '}'
        else:
            alignment = ''

        string += mtype + '}' + alignment
        string += '\n'

        shape = self.matrix.shape

        for (y, x), value in self._np.ndenumerate(self.matrix):
            if x:
                string += '&'
            string += str(value)

            if x == shape[1] - 1 and y != shape[0] - 1:
                string += r'\\' + '\n'

        string += '\n'

        string += r'\end{' + mtype + '}'

        super().dumps()

        return string
<commit_msg>Make Math.dumps code simpler<commit_after># -*- coding: utf-8 -*-
"""
This module implements the classes that deal with math.

..  :copyright: (c) 2014 by Jelte Fennema.
    :license: MIT, see License for more details.
"""

from .base_classes import LatexObject, Command, Container
from pylatex.package import Package

class Math(Container):

    """A class representing a math environment."""

    def __init__(self, data=None, inline=False):
        """.

        Args
        ----
        data: list
            Content of the math container.
        inline: bool
            If the math should be displayed inline or not.
        """

        self.inline = inline
        super().__init__(data)

    def dumps(self):
        """Return a LaTeX formatted string representing the object.

        Returns
        -------
        str

        """

        if self.inline:
            return '$' + super().dumps(token=' ') + '$'
        return '$$' + super().dumps(token=' ') + '$$\n'

class VectorName(Command):

    """A class representing a named vector."""

    def __init__(self, name):
        """.

        Args
        ----
        name: str
            Name of the vector
        """

        super().__init__('mathbf', arguments=name)

class Matrix(LatexObject):

    """A class representing a matrix."""

    # TODO: Convert this to an environment

    def __init__(self, matrix, mtype='p', alignment=None):
        r""".

        Args
        ----
        matrix: `numpy.ndarray` instance
            The matrix to display
        mtype: str
            What kind of brackets are used around the matrix. The different
            options and their corresponding brackets are:
            p = ( ), b = [ ], B = { }, v = \| \|, V = \|\| \|\|
        alignment: str
            How to align the content of the cells in the matrix. This is ``c``
            by default.

        References
        ----------
        * https://en.wikibooks.org/wiki/LaTeX/Mathematics#Matrices_and_arrays
        """

        import numpy as np
        self._np = np

        self.mtype = mtype
        self.matrix = matrix
        self.alignment = alignment

        super().__init__(packages=[Package('amsmath')])

    def dumps(self):
        """Return a string representin the matrix in LaTeX syntax.

        Returns
        -------
        str
        """

        string = r'\begin{'
        mtype = self.mtype + 'matrix'

        if self.alignment is not None:
            mtype += '*'
            alignment = '{' + self.alignment + '}'
        else:
            alignment = ''

        string += mtype + '}' + alignment
        string += '\n'

        shape = self.matrix.shape

        for (y, x), value in self._np.ndenumerate(self.matrix):
            if x:
                string += '&'
            string += str(value)

            if x == shape[1] - 1 and y != shape[0] - 1:
                string += r'\\' + '\n'

        string += '\n'

        string += r'\end{' + mtype + '}'

        super().dumps()

        return string
<|endoftext|>


<commit_before>from __future__ import division

import ctypes, os, shutil
import numpy as np

# constants
c = 299792458
eps0 = 8.854187817e-12
e = 1.602176565e-19
m_e = 9.10938291e-31
hbar = 1.054571726e-34
h = 2*np.pi*hbar
a0 = 4*np.pi*eps0*hbar**2/m_e/e**2
Ry = h*c * m_e*e**4/8/eps0**2/h**3/c

# find and import shared object / DLL
rootdir = os.path.dirname(os.path.realpath(__file__))

if os.name=="posix":
  lewenstein_so = ctypes.CDLL(os.path.join(rootdir,'lewenstein.so'))
elif os.name=="nt":
  bits = ctypes.sizeof(ctypes.c_voidp)*8
  archdirectory = os.path.join(rootdir,'dll' + str(bits))

  rootdirdll = os.path.join(rootdir,'lewenstein.dll')
  correctdll = os.path.join(archdirectory,'lewenstein.dll')
  if not os.path.exists(rootdirdll) or os.path.getsize(rootdirdll)!=os.path.getsize(correctdll):
    for filename in os.listdir(archdirectory): # for dependencies
      shutil.copy(os.path.join(archdirectory,filename), rootdir)

  lewenstein_so = ctypes.CDLL(os.path.join(rootdir,'lewenstein'))

# Note: explicitly setting lewenstein_so.*.argtypes/restype is necessary to prevent segfault on 64 bit:
# http://stackoverflow.com/questions/17240621/wrapping-simple-c-example-with-ctypes-segmentation-fault

# base class for dipole elements
class dipole_elements(object):
  dims = None
  pointer = None

  def __del__(self):
    raise NotImplementedError('override in subclasses')

# wrap H dipole elements
lewenstein_so.dipole_elements_H_double.argtypes = [ctypes.c_int, ctypes.c_double]
lewenstein_so.dipole_elements_H_double.restype = ctypes.c_void_p
lewenstein_so.dipole_elements_H_double_destroy.argtypes = [ctypes.c_int, ctypes.c_void_p]
lewenstein_so.dipole_elements_H_double_destroy.restype = None

class dipole_elements_H(dipole_elements):
  def __init__(self, dims, ip, wavelength=None):
    if wavelength:
      ip = sau_convert(ip, 'U', 'SAU', wavelength)

    alpha = 2*ip

    self.dims = dims
    self.pointer = lewenstein_so.dipole_elements_H_double(dims, alpha)

  def __del__(self):
    lewenstein_so.dipole_elements_H_double_destroy(self.dims, self.pointer)

# wrap symmetric interpolated dipole elements
lewenstein_so.dipole_elements_symmetric_interpolate_double.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_double, ctypes.c_void_p, ctypes.c_void_p]
lewenstein_so.dipole_elements_symmetric_interpolate_double.restype = ctypes.c_void_p
lewenstein_so.dipole_elements_symmetric_interpolate_double_destroy.argtypes = [ctypes.c_int, ctypes.c_void_p]
lewenstein_so.dipole_elements_symmetric_interpolate_double_destroy.restype = None

class dipole_elements_symmetric_interpolate(dipole_elements):
  # TODO: untested!

  _dr = None
  _di = None

  def __init__(self, dims, p, d, wavelength=None):
     if wavelength is not None:
       p = sau_convert(p, 'p', 'SAU', wavelength)
       d = sau_convert(d, 'd', 'SAU', wavelength)

     self.dims = dims

     N = p.size
     assert d.size==N

     dp = np.min(np.diff(p))
     assert np.isclose(dp, np.max(np.diff(p)), atol=0)

     # dr and di must not be garbage collected until destructor is called, so make it a property of this class
     self._dr = np.require(np.copy(d.real), np.double, ['C', 'A'])
     self._di = np.require(np.copy(d.imag), np.double, ['C', 'A'])

     self.pointer = lewenstein_so.dipole_elements_symmetric_interpolate_double(dims, N, dp, self._dr.ctypes.data, self._di.ctypes.data)

  def __del__(self):
     lewenstein_so.dipole_elements_symmetric_interpolate_double_destroy(self.dims, self.pointer)

# helper function to generate weights
def get_weights(tau,T=2*np.pi,periods_one=1,periods_soft=.5):
  interval_points = sum(tau-tau[0]<=periods_one*T)
  window_points = sum(tau-tau[0]<=periods_soft*T)

  r = np.ones(interval_points+window_points)
  r[-window_points:] = np.cos(np.pi/2 * np.arange(window_points)/window_points)**2

  return r

# helper function for unit conversion
def sau_convert(value, quantity, target, wavelength):
  # scaled atomic unit quantities expressed in SI units
  unit = {}
  unit['t'] = wavelength / c / (2*np.pi)
  unit['omega'] = 1/unit['t']
  unit['U'] = hbar * unit['omega'] # hbar*omega
  unit['q'] = e
  unit['s'] = a0 * np.sqrt(2*Ry/unit['U']) # [concluded from (23) in Lewenstein paper]
  unit['E'] = unit['U'] / unit['q'] / unit['s'] # [concluded from (2) in Lewenstein paper]
  unit['d'] = unit['q'] * unit['s'] # dipole element
  unit['m'] = m_e
  unit['p'] = unit['m'] * unit['s']/unit['t'] # momentum

  if target=="SI":
    return value * unit[quantity]
  elif target=="SAU":
    return value / unit[quantity]
  else:
    raise ValueError('target must be SI or SAU')

# wrap lewenstein function
lewenstein_so.lewenstein_double.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_double, ctypes.c_double, ctypes.c_void_p, ctypes.c_void_p]
lewenstein_so.lewenstein_double.restype = None

def lewenstein(t,Et,ip,wavelength=None,weights=None,at=None,dipole_elements=None,epsilon_t=1e-4):
  # default value for weights
  if weights is None and wavelength is None:
    weights = get_weights(t)
  elif weights is None and wavelength is not None:
    weights = get_weights(t, wavelength/c)

  # unit conversion
  if wavelength is not None:
    t = sau_convert(t, 't', 'SAU', wavelength)
    Et = sau_convert(Et, 'E', 'SAU', wavelength)
    ip = sau_convert(ip, 'U', 'SAU', wavelength)

  # default value for ground state amplitude
  if at is None: at = np.ones_like(t)

  # allocate memory for output
  output = np.empty_like(Et)

  # make sure t axis starts at zero
  t = t - t[0]

  # make sure we have appropriate memory layout before passing to C code
  t = np.require(t, np.double, ['C', 'A'])
  Et = np.require(Et, np.double, ['C', 'A'])
  weights = np.require(weights, np.double, ['C', 'A'])
  at = np.require(at, np.double, ['C', 'A'])
  output = np.require(output, np.double, ['C', 'A', 'W'])

  # get dimensions
  N = t.size
  dims = Et.shape[1] if len(Et.shape)>1 else 1
  weights_length = weights.size

  # check dimensions
  assert at.size==N
  assert Et.shape[0]==N
  assert dims in [1,2,3]
  assert Et.size==N*dims

  # default value for dipole elements
  if dipole_elements is None: dipole_elements = dipole_elements_H(dims, ip=ip)

  # call C function
  assert dipole_elements.dims==dims
  lewenstein_so.lewenstein_double(dims, N, t.ctypes.data, Et.ctypes.data, weights_length, weights.ctypes.data, at.ctypes.data, ip, epsilon_t, dipole_elements.pointer, output.ctypes.data)

  # unit conversion
  if wavelength is not None:
    output = sau_convert(output, 'd', 'SI', wavelength)

  return output

# wrap yakovlev function
lewenstein_so.yakovlev_double.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_double, ctypes.c_int, ctypes.c_int, ctypes.c_void_p]
lewenstein_so.yakovlev_double.restype = None

def yakovlev(t,Et,ip,at,wavelength=None,trajectories=2,skip_trajectories=0,max_periods=2,tb_window=None):
  # unit conversion
  if wavelength is not None:
    t = sau_convert(t, 't', 'SAU', wavelength)
    Et = sau_convert(Et, 'E', 'SAU', wavelength)
    ip = sau_convert(ip, 'U', 'SAU', wavelength)

  # default value for tb window
  if tb_window is None: tb_window = np.ones_like(t)

  # allocate memory for output
  output = np.empty_like(Et)

  # make sure t axis starts at zero
  t = t - t[0]

  # replace infinite max_periods by sensible value
  if max_periods is None or not np.isfinite(max_periods): max_periods = np.max(t)/2/np.pi + 1

  # make sure we have appropriate memory layout before passing to C code
  t = np.require(t, np.double, ['C', 'A'])
  Et = np.require(Et, np.double, ['C', 'A'])
  at = np.require(at, np.double, ['C', 'A'])
  tb_window = np.require(tb_window, np.double, ['C', 'A'])
  output = np.require(output, np.double, ['C', 'A', 'W'])

  # get dimensions
  N = t.size
  dims = Et.shape[1] if len(Et.shape)>1 else 1
  max_tau_i = np.where(t<2*np.pi*max_periods)[0][-1]

  # check dimensions
  assert at.size==N
  assert tb_window.size==N
  assert Et.shape[0]==N
  assert dims==1
  assert Et.size==N*dims

  # call C function
  lewenstein_so.yakovlev_double(dims, N, t.ctypes.data, Et.ctypes.data, max_tau_i, at.ctypes.data, tb_window.ctypes.data, ip, int(trajectories), int(skip_trajectories), output.ctypes.data)

  # unit conversion
  if wavelength is not None:
    output = sau_convert(output, 'd', 'SI', wavelength)

  return output

def cutoff(amplitude,ip,wavelength):
  """ expects SI units """

  # unit conversion
  amplitude = sau_convert(amplitude, 'E', 'SAU', wavelength)
  ip = sau_convert(ip, 'U', 'SAU', wavelength)
  omega = 2*np.pi/sau_convert(wavelength/c, 't', 'SAU', wavelength)

  up = amplitude**2/4/omega**2
  harmonic_order = ip + 3.17*up

  return harmonic_order

if __name__=="__main__":
  import pylab

  # compare to reference, in scaled atomic units
  n = 5

  t = np.arange(n)
  Et = t
  weights = np.ones_like(t)
  ip = 1

  d = lewenstein(t,Et,ip,None,weights)
  reference_d = [0.00000, 0.00000, -2.18318, -2.95464, -1.23753] # computed by Matlab/octave module
  assert np.allclose(d, reference_d, atol=1e-4)
  print "Test passed"

  # plot dipole response for pulse (using SI units)
  wavelength = 1000e-9
  T = wavelength/c
  t = np.linspace(-20*T,20*T,200*40+1)

  fwhm = 30e-15
  tau = fwhm/2/np.sqrt(np.log(np.sqrt(2)))
  Et = np.exp(-(t/tau)**2) * np.cos(2*np.pi/T*t) * np.sqrt(2*1e18/c/eps0)

  ip = 12.13*e # Xe

  d = lewenstein(t,Et,ip,wavelength)

  pylab.semilogy(np.fft.fftfreq(len(t), t[1]-t[0])/(1/T), abs(np.fft.fft(d))**2)
  pylab.xlim((0,100))
  pylab.show()
<commit_msg>test interpolated dipole elements in pylewenstein.py<commit_after>from __future__ import division

import ctypes, os, shutil
import numpy as np

# constants
c = 299792458
eps0 = 8.854187817e-12
e = 1.602176565e-19
m_e = 9.10938291e-31
hbar = 1.054571726e-34
h = 2*np.pi*hbar
a0 = 4*np.pi*eps0*hbar**2/m_e/e**2
Ry = h*c * m_e*e**4/8/eps0**2/h**3/c

# find and import shared object / DLL
rootdir = os.path.dirname(os.path.realpath(__file__))

if os.name=="posix":
  lewenstein_so = ctypes.CDLL(os.path.join(rootdir,'lewenstein.so'))
elif os.name=="nt":
  bits = ctypes.sizeof(ctypes.c_voidp)*8
  archdirectory = os.path.join(rootdir,'dll' + str(bits))

  rootdirdll = os.path.join(rootdir,'lewenstein.dll')
  correctdll = os.path.join(archdirectory,'lewenstein.dll')
  if not os.path.exists(rootdirdll) or os.path.getsize(rootdirdll)!=os.path.getsize(correctdll):
    for filename in os.listdir(archdirectory): # for dependencies
      shutil.copy(os.path.join(archdirectory,filename), rootdir)

  lewenstein_so = ctypes.CDLL(os.path.join(rootdir,'lewenstein'))

# Note: explicitly setting lewenstein_so.*.argtypes/restype is necessary to prevent segfault on 64 bit:
# http://stackoverflow.com/questions/17240621/wrapping-simple-c-example-with-ctypes-segmentation-fault

# base class for dipole elements
class dipole_elements(object):
  dims = None
  pointer = None

  def __del__(self):
    raise NotImplementedError('override in subclasses')

# wrap H dipole elements
lewenstein_so.dipole_elements_H_double.argtypes = [ctypes.c_int, ctypes.c_double]
lewenstein_so.dipole_elements_H_double.restype = ctypes.c_void_p
lewenstein_so.dipole_elements_H_double_destroy.argtypes = [ctypes.c_int, ctypes.c_void_p]
lewenstein_so.dipole_elements_H_double_destroy.restype = None

class dipole_elements_H(dipole_elements):
  def __init__(self, dims, ip, wavelength=None):
    if wavelength:
      ip = sau_convert(ip, 'U', 'SAU', wavelength)

    alpha = 2*ip

    self.dims = dims
    self.pointer = lewenstein_so.dipole_elements_H_double(dims, alpha)

  def __del__(self):
    lewenstein_so.dipole_elements_H_double_destroy(self.dims, self.pointer)

# wrap symmetric interpolated dipole elements
lewenstein_so.dipole_elements_symmetric_interpolate_double.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_double, ctypes.c_void_p, ctypes.c_void_p]
lewenstein_so.dipole_elements_symmetric_interpolate_double.restype = ctypes.c_void_p
lewenstein_so.dipole_elements_symmetric_interpolate_double_destroy.argtypes = [ctypes.c_int, ctypes.c_void_p]
lewenstein_so.dipole_elements_symmetric_interpolate_double_destroy.restype = None

class dipole_elements_symmetric_interpolate(dipole_elements):
  _dr = None
  _di = None

  def __init__(self, dims, p, d, wavelength=None):
     if wavelength is not None:
       p = sau_convert(p, 'p', 'SAU', wavelength)
       d = sau_convert(d, 'd', 'SAU', wavelength)

     self.dims = dims

     N = p.size
     assert d.size==N

     dp = np.min(np.diff(p))
     assert np.isclose(dp, np.max(np.diff(p)), atol=0)

     # dr and di must not be garbage collected until destructor is called, so make it a property of this class
     self._dr = np.require(np.copy(d.real), np.double, ['C', 'A'])
     self._di = np.require(np.copy(d.imag), np.double, ['C', 'A'])

     self.pointer = lewenstein_so.dipole_elements_symmetric_interpolate_double(dims, N, dp, self._dr.ctypes.data, self._di.ctypes.data)

  def __del__(self):
     lewenstein_so.dipole_elements_symmetric_interpolate_double_destroy(self.dims, self.pointer)

# helper function to generate weights
def get_weights(tau,T=2*np.pi,periods_one=1,periods_soft=.5):
  interval_points = sum(tau-tau[0]<=periods_one*T)
  window_points = sum(tau-tau[0]<=periods_soft*T)

  r = np.ones(interval_points+window_points)
  r[-window_points:] = np.cos(np.pi/2 * np.arange(window_points)/window_points)**2

  return r

# helper function for unit conversion
def sau_convert(value, quantity, target, wavelength):
  # scaled atomic unit quantities expressed in SI units
  unit = {}
  unit['t'] = wavelength / c / (2*np.pi)
  unit['omega'] = 1/unit['t']
  unit['U'] = hbar * unit['omega'] # hbar*omega
  unit['q'] = e
  unit['s'] = a0 * np.sqrt(2*Ry/unit['U']) # [concluded from (23) in Lewenstein paper]
  unit['E'] = unit['U'] / unit['q'] / unit['s'] # [concluded from (2) in Lewenstein paper]
  unit['d'] = unit['q'] * unit['s'] # dipole element
  unit['m'] = m_e
  unit['p'] = unit['m'] * unit['s']/unit['t'] # momentum

  if target=="SI":
    return value * unit[quantity]
  elif target=="SAU":
    return value / unit[quantity]
  else:
    raise ValueError('target must be SI or SAU')

# wrap lewenstein function
lewenstein_so.lewenstein_double.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_double, ctypes.c_double, ctypes.c_void_p, ctypes.c_void_p]
lewenstein_so.lewenstein_double.restype = None

def lewenstein(t,Et,ip,wavelength=None,weights=None,at=None,dipole_elements=None,epsilon_t=1e-4):
  # default value for weights
  if weights is None and wavelength is None:
    weights = get_weights(t)
  elif weights is None and wavelength is not None:
    weights = get_weights(t, wavelength/c)

  # unit conversion
  if wavelength is not None:
    t = sau_convert(t, 't', 'SAU', wavelength)
    Et = sau_convert(Et, 'E', 'SAU', wavelength)
    ip = sau_convert(ip, 'U', 'SAU', wavelength)

  # default value for ground state amplitude
  if at is None: at = np.ones_like(t)

  # allocate memory for output
  output = np.empty_like(Et)

  # make sure t axis starts at zero
  t = t - t[0]

  # make sure we have appropriate memory layout before passing to C code
  t = np.require(t, np.double, ['C', 'A'])
  Et = np.require(Et, np.double, ['C', 'A'])
  weights = np.require(weights, np.double, ['C', 'A'])
  at = np.require(at, np.double, ['C', 'A'])
  output = np.require(output, np.double, ['C', 'A', 'W'])

  # get dimensions
  N = t.size
  dims = Et.shape[1] if len(Et.shape)>1 else 1
  weights_length = weights.size

  # check dimensions
  assert at.size==N
  assert Et.shape[0]==N
  assert dims in [1,2,3]
  assert Et.size==N*dims

  # default value for dipole elements
  if dipole_elements is None: dipole_elements = dipole_elements_H(dims, ip=ip)

  # call C function
  assert dipole_elements.dims==dims
  lewenstein_so.lewenstein_double(dims, N, t.ctypes.data, Et.ctypes.data, weights_length, weights.ctypes.data, at.ctypes.data, ip, epsilon_t, dipole_elements.pointer, output.ctypes.data)

  # unit conversion
  if wavelength is not None:
    output = sau_convert(output, 'd', 'SI', wavelength)

  return output

# wrap yakovlev function
lewenstein_so.yakovlev_double.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_double, ctypes.c_int, ctypes.c_int, ctypes.c_void_p]
lewenstein_so.yakovlev_double.restype = None

def yakovlev(t,Et,ip,at,wavelength=None,trajectories=2,skip_trajectories=0,max_periods=2,tb_window=None):
  # unit conversion
  if wavelength is not None:
    t = sau_convert(t, 't', 'SAU', wavelength)
    Et = sau_convert(Et, 'E', 'SAU', wavelength)
    ip = sau_convert(ip, 'U', 'SAU', wavelength)

  # default value for tb window
  if tb_window is None: tb_window = np.ones_like(t)

  # allocate memory for output
  output = np.empty_like(Et)

  # make sure t axis starts at zero
  t = t - t[0]

  # replace infinite max_periods by sensible value
  if max_periods is None or not np.isfinite(max_periods): max_periods = np.max(t)/2/np.pi + 1

  # make sure we have appropriate memory layout before passing to C code
  t = np.require(t, np.double, ['C', 'A'])
  Et = np.require(Et, np.double, ['C', 'A'])
  at = np.require(at, np.double, ['C', 'A'])
  tb_window = np.require(tb_window, np.double, ['C', 'A'])
  output = np.require(output, np.double, ['C', 'A', 'W'])

  # get dimensions
  N = t.size
  dims = Et.shape[1] if len(Et.shape)>1 else 1
  max_tau_i = np.where(t<2*np.pi*max_periods)[0][-1]

  # check dimensions
  assert at.size==N
  assert tb_window.size==N
  assert Et.shape[0]==N
  assert dims==1
  assert Et.size==N*dims

  # call C function
  lewenstein_so.yakovlev_double(dims, N, t.ctypes.data, Et.ctypes.data, max_tau_i, at.ctypes.data, tb_window.ctypes.data, ip, int(trajectories), int(skip_trajectories), output.ctypes.data)

  # unit conversion
  if wavelength is not None:
    output = sau_convert(output, 'd', 'SI', wavelength)

  return output

def cutoff(amplitude,ip,wavelength):
  """ expects SI units """

  # unit conversion
  amplitude = sau_convert(amplitude, 'E', 'SAU', wavelength)
  ip = sau_convert(ip, 'U', 'SAU', wavelength)
  omega = 2*np.pi/sau_convert(wavelength/c, 't', 'SAU', wavelength)

  up = amplitude**2/4/omega**2
  harmonic_order = ip + 3.17*up

  return harmonic_order

if __name__=="__main__":
  import pylab

  # compare to reference, in scaled atomic units
  n = 5

  t = np.arange(n)
  Et = t
  weights = np.ones_like(t)
  ip = 1

  d = lewenstein(t,Et,ip,None,weights)
  reference_d = [0.00000, 0.00000, -2.18318, -2.95464, -1.23753] # computed by Matlab/octave module
  assert np.allclose(d, reference_d, atol=1e-4)
  print "Test passed"

  # plot dipole response for pulse (using SI units)
  wavelength = 1000e-9
  T = wavelength/c
  t = np.linspace(-20*T,20*T,200*40+1)

  fwhm = 30e-15
  tau = fwhm/2/np.sqrt(np.log(np.sqrt(2)))
  Et = np.exp(-(t/tau)**2) * np.cos(2*np.pi/T*t) * np.sqrt(2*1e18/c/eps0)

  ip = 12.13*e # Xe

  d = lewenstein(t,Et,ip,wavelength)

  Dp = np.arange(0,10,1e-4)
  alpha = 2*sau_convert(ip,'U','SAU',wavelength)
  Dd = 1j*(2**(7/2)*alpha**(5/4)/np.pi) * Dp / (Dp**2+alpha)**3
  dint = dipole_elements_symmetric_interpolate(1, Dp, Dd)
  d2 = lewenstein(t,Et,ip,wavelength,dipole_elements=dint)

  pylab.semilogy(np.fft.fftfreq(len(t), t[1]-t[0])/(1/T), abs(np.fft.fft(d))**2)
  pylab.semilogy(np.fft.fftfreq(len(t), t[1]-t[0])/(1/T), abs(np.fft.fft(d2))**2, label='interpolated d')
  pylab.legend()
  pylab.xlim((0,100))
  pylab.show()
<|endoftext|>


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

"""Tools for dealing with Mongo's BSON data representation.

Generally not needed to be used by application developers."""

import types
import struct
import random
import re
import datetime
import calendar
import logging

from binary import Binary
from code import Code
from objectid import ObjectId
from dbref import DBRef
from son import SON
from errors import InvalidBSON, InvalidDocument, UnsupportedTag

try:
    import _cbson
    _use_c = True
except ImportError:
    _use_c = False

_logger = logging.getLogger("pymongo.bson")
# _logger.setLevel(logging.DEBUG)
# _logger.addHandler(logging.StreamHandler())

def _get_int(data):
    try:
        value = struct.unpack("<i", data[:4])[0]
    except struct.error:
        raise InvalidBSON()

    return (value, data[4:])

def _get_c_string(data):
    try:
        end = data.index("\x00")
    except ValueError:
        raise InvalidBSON()

    return (unicode(data[:end], "utf-8"), data[end + 1:])

def _make_c_string(string, encode=True):
    if encode and isinstance(string, unicode):
        return string.encode("utf-8") + "\x00"
    return string + "\x00"

def _validate_number(data):
    assert len(data) >= 8
    return data[8:]

def _validate_string(data):
    (length, data) = _get_int(data)
    assert len(data) >= length
    assert data[length - 1] == "\x00"
    return data[length:]

def _validate_object(data):
    return _validate_document(data, None)

_valid_array_name = re.compile("^\d+$")
def _validate_array(data):
    return _validate_document(data, _valid_array_name)

def _validate_binary(data):
    (length, data) = _get_int(data)
    # + 1 for the subtype byte
    assert len(data) >= length + 1
    return data[length + 1:]

def _validate_undefined(data):
    return data

_OID_SIZE = 12
def _validate_oid(data):
    assert len(data) >= _OID_SIZE
    return data[_OID_SIZE:]

def _validate_boolean(data):
    assert len(data) >= 1
    return data[1:]

_DATE_SIZE = 8
def _validate_date(data):
    assert len(data) >= _DATE_SIZE
    return data[_DATE_SIZE:]

_validate_null = _validate_undefined

def _validate_regex(data):
    (regex, data) = _get_c_string(data)
    (options, data) = _get_c_string(data)
    return data

def _validate_ref(data):
    data = _validate_string(data)
    return _validate_oid(data)

_validate_code = _validate_string

# still not sure what is actually stored here, but i know how big it is...
def _validate_code_w_scope(data):
    (length, data) = _get_int(data)
    assert len(data) >= length + 1
    return data[length + 1:]

_validate_symbol = _validate_string

def _validate_number_int(data):
    assert len(data) >= 4
    return data[4:]

_element_validator = {
    "\x01": _validate_number,
    "\x02": _validate_string,
    "\x03": _validate_object,
    "\x04": _validate_array,
    "\x05": _validate_binary,
    "\x06": _validate_undefined,
    "\x07": _validate_oid,
    "\x08": _validate_boolean,
    "\x09": _validate_date,
    "\x0A": _validate_null,
    "\x0B": _validate_regex,
    "\x0C": _validate_ref,
    "\x0D": _validate_code,
    "\x0E": _validate_symbol,
    "\x0F": _validate_code_w_scope,
    "\x10": _validate_number_int,
}

def _validate_element_data(type, data):
    try:
        return _element_validator[type](data)
    except KeyError:
        raise InvalidBSON("unrecognized type: %s" % type)

def _validate_element(data, valid_name):
    element_type = data[0]
    (element_name, data) = _get_c_string(data[1:])
    if valid_name:
        assert valid_name.match(element_name), "%r doesn't match %s" % (element_name, valid_name.pattern)
    return _validate_element_data(element_type, data)

def _validate_elements(data, valid_name):
    while data:
        data = _validate_element(data, valid_name)

def _validate_document(data, valid_name=None):
    try:
        obj_size = struct.unpack("<i", data[:4])[0]
    except struct.error:
        raise InvalidBSON()

    assert obj_size <= len(data)
    obj = data[4:obj_size]
    assert len(obj)

    eoo = obj[-1]
    assert eoo == "\x00"

    elements = obj[:-1]
    _validate_elements(elements, valid_name)

    return data[obj_size:]

def _get_number(data):
    return (struct.unpack("<d", data[:8])[0], data[8:])

def _get_string(data):
    return _get_c_string(data[4:])

def _get_object(data):
    return _bson_to_dict(data)

def _get_array(data):
    (obj, data) = _get_object(data)
    result = []
    i = 0
    while True:
        try:
            result.append(obj[str(i)])
            i += 1
        except KeyError:
            break
    return (result, data)

def _get_binary(data):
    (length, data) = _get_int(data)
    subtype = ord(data[0])
    data = data[1:]
    if subtype == 2:
        (length2, data) = _get_int(data)
        if length2 != length - 4:
            raise InvalidBSON("invalid binary (st 2) - lengths don't match!")
        length = length2
    return (Binary(data[:length], subtype), data[length:])

def _get_oid(data):
    oid = _shuffle_oid(data[:12])
    return (ObjectId(oid), data[12:])

def _get_boolean(data):
    return (data[0] == "\x01", data[1:])

def _get_date(data):
    seconds = float(struct.unpack("<q", data[:8])[0]) / 1000.0
    return (datetime.datetime.utcfromtimestamp(seconds), data[8:])

def _get_null(data):
    return (None, data)

def _get_regex(data):
    (pattern, data) = _get_c_string(data)
    (bson_flags, data) = _get_c_string(data)
    flags = 0
    if "i" in bson_flags:
        flags |= re.IGNORECASE
    if "l" in bson_flags:
        flags |= re.LOCALE
    if "m" in bson_flags:
        flags |= re.MULTILINE
    if "s" in bson_flags:
        flags |= re.DOTALL
    if "u" in bson_flags:
        flags |= re.UNICODE
    if "x" in bson_flags:
        flags |= re.VERBOSE
    return (re.compile(pattern, flags), data)

def _get_ref(data):
    (collection, data) = _get_c_string(data[4:])
    (oid, data) = _get_oid(data)
    return (DBRef(collection, oid), data)

_element_getter = {
    "\x01": _get_number,
    "\x02": _get_string,
    "\x03": _get_object,
    "\x04": _get_array,
    "\x05": _get_binary,
    "\x06": _get_null, # undefined
    "\x07": _get_oid,
    "\x08": _get_boolean,
    "\x09": _get_date,
    "\x0A": _get_null,
    "\x0B": _get_regex,
    "\x0C": _get_ref,
    "\x0D": _get_string, # code
    "\x0E": _get_string, # symbol
#    "\x0F": _get_code_w_scope
    "\x10": _get_int, # number_int
}

def _element_to_dict(data):
    element_type = data[0]
    (element_name, data) = _get_c_string(data[1:])
    (value, data) = _element_getter[element_type](data)
    return (element_name, value, data)

def _elements_to_dict(data):
    result = SON()
    while data:
        (key, value, data) = _element_to_dict(data)
        result[key] = value
    return result

def _bson_to_dict(data):
    obj_size = struct.unpack("<i", data[:4])[0]
    elements = data[4:obj_size - 1]
    return (_elements_to_dict(elements), data[obj_size:])
if _use_c:
    _bson_to_dict = _cbson._bson_to_dict

def _shuffle_oid(data):
    return data[7::-1] + data[:7:-1]

_RE_TYPE = type(_valid_array_name)
def _element_to_bson(key, value):
    name = _make_c_string(key)
    if isinstance(value, float):
        return "\x01" + name + struct.pack("<d", value)
    if isinstance(value, Binary):
        subtype = value.subtype()
        if subtype == 2:
            value = struct.pack("<i", len(value)) + value
        return "\x05" + name + struct.pack("<i", len(value)) + chr(subtype) + value
    if isinstance(value, Code):
        cstring = _make_c_string(value, False)
        length = struct.pack("<i", len(cstring))
        return "\x0D" + name + length + cstring
    if isinstance(value, str):
        cstring = _make_c_string(value, False)
        length = struct.pack("<i", len(cstring))
        return "\x02" + name + length + cstring
    if isinstance(value, unicode):
        cstring = _make_c_string(value)
        length = struct.pack("<i", len(cstring))
        return "\x02" + name + length + cstring
    if isinstance(value, dict):
        return "\x03" + name + _dict_to_bson(value)
    if isinstance(value, (list, tuple)):
        as_dict = SON(zip([str(i) for i in range(len(value))], value))
        return "\x04" + name + _dict_to_bson(as_dict)
    if isinstance(value, ObjectId):
        return "\x07" + name + _shuffle_oid(str(value))
    if value is True:
        return "\x08" + name + "\x01"
    if value is False:
        return "\x08" + name + "\x00"
    if isinstance(value, (int, long)):
        return "\x10" + name + struct.pack("<i", value)
    if isinstance(value, datetime.datetime):
        millis = int(calendar.timegm(value.timetuple()) * 1000 + value.microsecond / 1000)
        return "\x09" + name + struct.pack("<q", millis)
    if value is None:
        return "\x0A" + name
    if isinstance(value, _RE_TYPE):
        pattern = value.pattern
        flags = ""
        if value.flags & re.IGNORECASE:
            flags += "i"
        if value.flags & re.LOCALE:
            flags += "l"
        if value.flags & re.MULTILINE:
            flags += "m"
        if value.flags & re.DOTALL:
            flags += "s"
        if value.flags & re.UNICODE:
            flags += "u"
        if value.flags & re.VERBOSE:
            flags += "x"
        return "\x0B" + name + _make_c_string(pattern) + _make_c_string(flags, False)
    if isinstance(value, DBRef):
        ns = _make_c_string(value.collection())
        return "\x0C" + name + struct.pack("<i", len(ns)) + ns + _shuffle_oid(str(value.id()))
    raise InvalidDocument("cannot convert value of type %s to bson" % type(value))

def _dict_to_bson(dict):
    try:
        elements = "".join([_element_to_bson(key, value) for (key, value) in dict.iteritems()])
    except AttributeError:
        raise TypeError("argument to from_dict must be a mapping type")

    length = len(elements) + 5
    return struct.pack("<i", length) + elements + "\x00"
if _use_c:
    _dict_to_bson = _cbson._dict_to_bson

def _to_dicts(data):
    """Convert binary data to sequence of SON objects.

    Data must be concatenated strings of valid BSON data.

    :Parameters:
      - `data`: bson data
    """
    dicts = []
    while len(data):
        (son, data) = _bson_to_dict(data)
        dicts.append(son)
    return dicts

def _to_dict(data):
    (son, _) = _bson_to_dict(data)
    return son

def is_valid(bson):
    """Validate that the given string represents valid BSON data.

    Raises TypeError if the data is not an instance of a subclass of str.
    Returns True if the data represents a valid BSON object, False otherwise.

    :Parameters:
      - `bson`: the data to be validated
    """
    if not isinstance(bson, types.StringType):
        raise TypeError("BSON data must be an instance of a subclass of str")

    try:
        remainder = _validate_document(bson)
        return remainder == ""
    except (AssertionError, InvalidBSON):
        return False

class BSON(str):
    """BSON data.

    Represents binary data storable in and retrievable from Mongo.
    """
    def __new__(cls, bson):
        """Initialize a new BSON object with some data.

        Raises TypeError if `bson` is not an instance of str.

        :Parameters:
          - `bson`: the initial data
        """
        return str.__new__(cls, bson)

    @classmethod
    def from_dict(cls, dict):
        """Create a new BSON object from a python mapping type (like dict).

        Raises TypeError if the argument is not a mapping type, or contains keys
        that are not instance of (str, unicode). Raises InvalidDocument if the
        dictionary cannot be converted to BSON.

        :Parameters:
          - `dict`: mapping type representing a Mongo document
        """
        return cls(_dict_to_bson(dict))

    def to_dict(self):
        """Get the dictionary representation of this data."""
        (son, _) = _bson_to_dict(self)
        return son
<commit_msg>Python BSON decoder returns a plain dictionary instead of SON.<commit_after># Copyright 2009 10gen, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tools for dealing with Mongo's BSON data representation.

Generally not needed to be used by application developers."""

import types
import struct
import random
import re
import datetime
import calendar
import logging

from binary import Binary
from code import Code
from objectid import ObjectId
from dbref import DBRef
from son import SON
from errors import InvalidBSON, InvalidDocument, UnsupportedTag

try:
    import _cbson
    _use_c = True
except ImportError:
    _use_c = False

_logger = logging.getLogger("pymongo.bson")
# _logger.setLevel(logging.DEBUG)
# _logger.addHandler(logging.StreamHandler())

def _get_int(data):
    try:
        value = struct.unpack("<i", data[:4])[0]
    except struct.error:
        raise InvalidBSON()

    return (value, data[4:])

def _get_c_string(data):
    try:
        end = data.index("\x00")
    except ValueError:
        raise InvalidBSON()

    return (unicode(data[:end], "utf-8"), data[end + 1:])

def _make_c_string(string, encode=True):
    if encode and isinstance(string, unicode):
        return string.encode("utf-8") + "\x00"
    return string + "\x00"

def _validate_number(data):
    assert len(data) >= 8
    return data[8:]

def _validate_string(data):
    (length, data) = _get_int(data)
    assert len(data) >= length
    assert data[length - 1] == "\x00"
    return data[length:]

def _validate_object(data):
    return _validate_document(data, None)

_valid_array_name = re.compile("^\d+$")
def _validate_array(data):
    return _validate_document(data, _valid_array_name)

def _validate_binary(data):
    (length, data) = _get_int(data)
    # + 1 for the subtype byte
    assert len(data) >= length + 1
    return data[length + 1:]

def _validate_undefined(data):
    return data

_OID_SIZE = 12
def _validate_oid(data):
    assert len(data) >= _OID_SIZE
    return data[_OID_SIZE:]

def _validate_boolean(data):
    assert len(data) >= 1
    return data[1:]

_DATE_SIZE = 8
def _validate_date(data):
    assert len(data) >= _DATE_SIZE
    return data[_DATE_SIZE:]

_validate_null = _validate_undefined

def _validate_regex(data):
    (regex, data) = _get_c_string(data)
    (options, data) = _get_c_string(data)
    return data

def _validate_ref(data):
    data = _validate_string(data)
    return _validate_oid(data)

_validate_code = _validate_string

# still not sure what is actually stored here, but i know how big it is...
def _validate_code_w_scope(data):
    (length, data) = _get_int(data)
    assert len(data) >= length + 1
    return data[length + 1:]

_validate_symbol = _validate_string

def _validate_number_int(data):
    assert len(data) >= 4
    return data[4:]

_element_validator = {
    "\x01": _validate_number,
    "\x02": _validate_string,
    "\x03": _validate_object,
    "\x04": _validate_array,
    "\x05": _validate_binary,
    "\x06": _validate_undefined,
    "\x07": _validate_oid,
    "\x08": _validate_boolean,
    "\x09": _validate_date,
    "\x0A": _validate_null,
    "\x0B": _validate_regex,
    "\x0C": _validate_ref,
    "\x0D": _validate_code,
    "\x0E": _validate_symbol,
    "\x0F": _validate_code_w_scope,
    "\x10": _validate_number_int,
}

def _validate_element_data(type, data):
    try:
        return _element_validator[type](data)
    except KeyError:
        raise InvalidBSON("unrecognized type: %s" % type)

def _validate_element(data, valid_name):
    element_type = data[0]
    (element_name, data) = _get_c_string(data[1:])
    if valid_name:
        assert valid_name.match(element_name), "%r doesn't match %s" % (element_name, valid_name.pattern)
    return _validate_element_data(element_type, data)

def _validate_elements(data, valid_name):
    while data:
        data = _validate_element(data, valid_name)

def _validate_document(data, valid_name=None):
    try:
        obj_size = struct.unpack("<i", data[:4])[0]
    except struct.error:
        raise InvalidBSON()

    assert obj_size <= len(data)
    obj = data[4:obj_size]
    assert len(obj)

    eoo = obj[-1]
    assert eoo == "\x00"

    elements = obj[:-1]
    _validate_elements(elements, valid_name)

    return data[obj_size:]

def _get_number(data):
    return (struct.unpack("<d", data[:8])[0], data[8:])

def _get_string(data):
    return _get_c_string(data[4:])

def _get_object(data):
    return _bson_to_dict(data)

def _get_array(data):
    (obj, data) = _get_object(data)
    result = []
    i = 0
    while True:
        try:
            result.append(obj[str(i)])
            i += 1
        except KeyError:
            break
    return (result, data)

def _get_binary(data):
    (length, data) = _get_int(data)
    subtype = ord(data[0])
    data = data[1:]
    if subtype == 2:
        (length2, data) = _get_int(data)
        if length2 != length - 4:
            raise InvalidBSON("invalid binary (st 2) - lengths don't match!")
        length = length2
    return (Binary(data[:length], subtype), data[length:])

def _get_oid(data):
    oid = _shuffle_oid(data[:12])
    return (ObjectId(oid), data[12:])

def _get_boolean(data):
    return (data[0] == "\x01", data[1:])

def _get_date(data):
    seconds = float(struct.unpack("<q", data[:8])[0]) / 1000.0
    return (datetime.datetime.utcfromtimestamp(seconds), data[8:])

def _get_null(data):
    return (None, data)

def _get_regex(data):
    (pattern, data) = _get_c_string(data)
    (bson_flags, data) = _get_c_string(data)
    flags = 0
    if "i" in bson_flags:
        flags |= re.IGNORECASE
    if "l" in bson_flags:
        flags |= re.LOCALE
    if "m" in bson_flags:
        flags |= re.MULTILINE
    if "s" in bson_flags:
        flags |= re.DOTALL
    if "u" in bson_flags:
        flags |= re.UNICODE
    if "x" in bson_flags:
        flags |= re.VERBOSE
    return (re.compile(pattern, flags), data)

def _get_ref(data):
    (collection, data) = _get_c_string(data[4:])
    (oid, data) = _get_oid(data)
    return (DBRef(collection, oid), data)

_element_getter = {
    "\x01": _get_number,
    "\x02": _get_string,
    "\x03": _get_object,
    "\x04": _get_array,
    "\x05": _get_binary,
    "\x06": _get_null, # undefined
    "\x07": _get_oid,
    "\x08": _get_boolean,
    "\x09": _get_date,
    "\x0A": _get_null,
    "\x0B": _get_regex,
    "\x0C": _get_ref,
    "\x0D": _get_string, # code
    "\x0E": _get_string, # symbol
#    "\x0F": _get_code_w_scope
    "\x10": _get_int, # number_int
}

def _element_to_dict(data):
    element_type = data[0]
    (element_name, data) = _get_c_string(data[1:])
    (value, data) = _element_getter[element_type](data)
    return (element_name, value, data)

def _elements_to_dict(data):
    result = {}
    while data:
        (key, value, data) = _element_to_dict(data)
        result[key] = value
    return result

def _bson_to_dict(data):
    obj_size = struct.unpack("<i", data[:4])[0]
    elements = data[4:obj_size - 1]
    return (_elements_to_dict(elements), data[obj_size:])
if _use_c:
    _bson_to_dict = _cbson._bson_to_dict

def _shuffle_oid(data):
    return data[7::-1] + data[:7:-1]

_RE_TYPE = type(_valid_array_name)
def _element_to_bson(key, value):
    name = _make_c_string(key)
    if isinstance(value, float):
        return "\x01" + name + struct.pack("<d", value)
    if isinstance(value, Binary):
        subtype = value.subtype()
        if subtype == 2:
            value = struct.pack("<i", len(value)) + value
        return "\x05" + name + struct.pack("<i", len(value)) + chr(subtype) + value
    if isinstance(value, Code):
        cstring = _make_c_string(value, False)
        length = struct.pack("<i", len(cstring))
        return "\x0D" + name + length + cstring
    if isinstance(value, str):
        cstring = _make_c_string(value, False)
        length = struct.pack("<i", len(cstring))
        return "\x02" + name + length + cstring
    if isinstance(value, unicode):
        cstring = _make_c_string(value)
        length = struct.pack("<i", len(cstring))
        return "\x02" + name + length + cstring
    if isinstance(value, dict):
        return "\x03" + name + _dict_to_bson(value)
    if isinstance(value, (list, tuple)):
        as_dict = SON(zip([str(i) for i in range(len(value))], value))
        return "\x04" + name + _dict_to_bson(as_dict)
    if isinstance(value, ObjectId):
        return "\x07" + name + _shuffle_oid(str(value))
    if value is True:
        return "\x08" + name + "\x01"
    if value is False:
        return "\x08" + name + "\x00"
    if isinstance(value, (int, long)):
        return "\x10" + name + struct.pack("<i", value)
    if isinstance(value, datetime.datetime):
        millis = int(calendar.timegm(value.timetuple()) * 1000 + value.microsecond / 1000)
        return "\x09" + name + struct.pack("<q", millis)
    if value is None:
        return "\x0A" + name
    if isinstance(value, _RE_TYPE):
        pattern = value.pattern
        flags = ""
        if value.flags & re.IGNORECASE:
            flags += "i"
        if value.flags & re.LOCALE:
            flags += "l"
        if value.flags & re.MULTILINE:
            flags += "m"
        if value.flags & re.DOTALL:
            flags += "s"
        if value.flags & re.UNICODE:
            flags += "u"
        if value.flags & re.VERBOSE:
            flags += "x"
        return "\x0B" + name + _make_c_string(pattern) + _make_c_string(flags, False)
    if isinstance(value, DBRef):
        ns = _make_c_string(value.collection())
        return "\x0C" + name + struct.pack("<i", len(ns)) + ns + _shuffle_oid(str(value.id()))
    raise InvalidDocument("cannot convert value of type %s to bson" % type(value))

def _dict_to_bson(dict):
    try:
        elements = "".join([_element_to_bson(key, value) for (key, value) in dict.iteritems()])
    except AttributeError:
        raise TypeError("argument to from_dict must be a mapping type")

    length = len(elements) + 5
    return struct.pack("<i", length) + elements + "\x00"
if _use_c:
    _dict_to_bson = _cbson._dict_to_bson

def _to_dicts(data):
    """Convert binary data to sequence of SON objects.

    Data must be concatenated strings of valid BSON data.

    :Parameters:
      - `data`: bson data
    """
    dicts = []
    while len(data):
        (son, data) = _bson_to_dict(data)
        dicts.append(son)
    return dicts

def _to_dict(data):
    (son, _) = _bson_to_dict(data)
    return son

def is_valid(bson):
    """Validate that the given string represents valid BSON data.

    Raises TypeError if the data is not an instance of a subclass of str.
    Returns True if the data represents a valid BSON object, False otherwise.

    :Parameters:
      - `bson`: the data to be validated
    """
    if not isinstance(bson, types.StringType):
        raise TypeError("BSON data must be an instance of a subclass of str")

    try:
        remainder = _validate_document(bson)
        return remainder == ""
    except (AssertionError, InvalidBSON):
        return False

class BSON(str):
    """BSON data.

    Represents binary data storable in and retrievable from Mongo.
    """
    def __new__(cls, bson):
        """Initialize a new BSON object with some data.

        Raises TypeError if `bson` is not an instance of str.

        :Parameters:
          - `bson`: the initial data
        """
        return str.__new__(cls, bson)

    @classmethod
    def from_dict(cls, dict):
        """Create a new BSON object from a python mapping type (like dict).

        Raises TypeError if the argument is not a mapping type, or contains keys
        that are not instance of (str, unicode). Raises InvalidDocument if the
        dictionary cannot be converted to BSON.

        :Parameters:
          - `dict`: mapping type representing a Mongo document
        """
        return cls(_dict_to_bson(dict))

    def to_dict(self):
        """Get the dictionary representation of this data."""
        (son, _) = _bson_to_dict(self)
        return son
<|endoftext|>


<commit_before>"""Define MyQ devices."""
import asyncio
import logging
from datetime import datetime
from typing import TYPE_CHECKING, Optional, List

from .const import DEVICE_TYPE, WAIT_TIMEOUT
from .errors import RequestError, MyQError

if TYPE_CHECKING:
    from .api import API

_LOGGER = logging.getLogger(__name__)

class MyQDevice:
    """Define a generic device."""

    def __init__(
        self, api: "API", device_json: dict, account: str, state_update: datetime
    ) -> None:
        """Initialize.
        :type account: str
        """
        self._api = api  # type: "API"
        self._account = account  # type: str
        self.device_json = device_json  # type: dict
        self.state_update = state_update  # type: datetime
        self._device_state = None  # Type: Optional[str]

    @property
    def account(self) -> str:
        """Return account associated with device"""
        return self._account

    @property
    def device_family(self) -> str:
        """Return the family in which this device lives."""
        return self.device_json["device_family"]

    @property
    def device_id(self) -> str:
        """Return the device ID (serial number)."""
        return self.device_json["serial_number"]

    @property
    def device_platform(self) -> str:
        """Return the device platform."""
        return self.device_json["device_platform"]

    @property
    def device_type(self) -> str:
        """Return the device type."""
        return self.device_json[DEVICE_TYPE]

    @property
    def firmware_version(self) -> Optional[str]:
        """Return the family in which this device lives."""
        return self.device_json["state"].get("firmware_version")

    @property
    def name(self) -> bool:
        """Return the device name."""
        return self.device_json["name"]

    @property
    def online(self) -> bool:
        """Return whether the device is online."""
        return self.device_json["state"].get("online") is True

    @property
    def parent_device_id(self) -> Optional[str]:
        """Return the device ID (serial number) of this device's parent."""
        return self.device_json.get("parent_device_id")

    @property
    def href(self) -> Optional[str]:
        """Return the hyperlinks of the device."""
        return self.device_json.get("href")

    @property
    def state(self) -> Optional[str]:
        return self._device_state or self.device_state

    @state.setter
    def state(self, value: str) -> None:
        """Set the current state of the device."""
        self._device_state = value

    @property
    def device_state(self) -> Optional[str]:
        return None

    @property
    def close_allowed(self) -> bool:
        """Return whether the device can be closed unattended."""
        return False

    @property
    def open_allowed(self) -> bool:
        """Return whether the device can be opened unattended."""
        return False

    async def _send_state_command(self, url: str, command: str) -> None:
        """Instruct the API to change the state of the device."""
        # If the user tries to open or close, say, a gateway, throw an exception:
        if not self.state:
            raise RequestError(
                f"Cannot change state of device type: {self.device_type}"
            )

        _LOGGER.debug(f"Sending command {command} for {self.name}")
        await self._api.request(
            method="put",
            returns="response",
            url=url,
        )

    async def update(self) -> None:
        """Get the latest info for this device."""
        await self._api.update_device_info()

    async def wait_for_state(
        self,
        current_state: List,
        new_state: List,
        last_state_update: datetime,
        timeout: int = WAIT_TIMEOUT,
    ) -> bool:
        # First wait until door state is actually updated.
        _LOGGER.debug(f"Waiting until device state has been updated for {self.name}")
        wait_timeout = timeout
        while (
            last_state_update
            == self.device_json["state"].get("last_update", datetime.utcnow())
            and wait_timeout > 0
        ):
            wait_timeout = wait_timeout - 5
            await asyncio.sleep(5)
            try:
                await self._api.update_device_info(for_account=self.account)
            except MyQError:
                # Ignoring
                pass

        # Wait until the state is to what we want it to be
        _LOGGER.debug(f"Waiting until device state for {self.name} is {new_state}")
        wait_timeout = timeout
        while self.state in current_state and wait_timeout > 0:
            wait_timeout = wait_timeout - 5
            await asyncio.sleep(5)
            try:
                await self._api.update_device_info(for_account=self.account)
            except MyQError:
                # Ignoring
                pass

        # Reset self.state ensuring it reflects actual device state. Only do this if state is still what it would
        # have been, this to ensure if something else had updated it to something else we don't override.
        if self._device_state == current_state:
            self.state = None

        return self.state in new_state
<commit_msg>Fix resetting state<commit_after>"""Define MyQ devices."""
import asyncio
import logging
from datetime import datetime
from typing import TYPE_CHECKING, Optional, List

from .const import DEVICE_TYPE, WAIT_TIMEOUT
from .errors import RequestError, MyQError

if TYPE_CHECKING:
    from .api import API

_LOGGER = logging.getLogger(__name__)

class MyQDevice:
    """Define a generic device."""

    def __init__(
        self, api: "API", device_json: dict, account: str, state_update: datetime
    ) -> None:
        """Initialize.
        :type account: str
        """
        self._api = api  # type: "API"
        self._account = account  # type: str
        self.device_json = device_json  # type: dict
        self.state_update = state_update  # type: datetime
        self._device_state = None  # Type: Optional[str]

    @property
    def account(self) -> str:
        """Return account associated with device"""
        return self._account

    @property
    def device_family(self) -> str:
        """Return the family in which this device lives."""
        return self.device_json["device_family"]

    @property
    def device_id(self) -> str:
        """Return the device ID (serial number)."""
        return self.device_json["serial_number"]

    @property
    def device_platform(self) -> str:
        """Return the device platform."""
        return self.device_json["device_platform"]

    @property
    def device_type(self) -> str:
        """Return the device type."""
        return self.device_json[DEVICE_TYPE]

    @property
    def firmware_version(self) -> Optional[str]:
        """Return the family in which this device lives."""
        return self.device_json["state"].get("firmware_version")

    @property
    def name(self) -> bool:
        """Return the device name."""
        return self.device_json["name"]

    @property
    def online(self) -> bool:
        """Return whether the device is online."""
        return self.device_json["state"].get("online") is True

    @property
    def parent_device_id(self) -> Optional[str]:
        """Return the device ID (serial number) of this device's parent."""
        return self.device_json.get("parent_device_id")

    @property
    def href(self) -> Optional[str]:
        """Return the hyperlinks of the device."""
        return self.device_json.get("href")

    @property
    def state(self) -> Optional[str]:
        return self._device_state or self.device_state

    @state.setter
    def state(self, value: str) -> None:
        """Set the current state of the device."""
        self._device_state = value

    @property
    def device_state(self) -> Optional[str]:
        return None

    @property
    def close_allowed(self) -> bool:
        """Return whether the device can be closed unattended."""
        return False

    @property
    def open_allowed(self) -> bool:
        """Return whether the device can be opened unattended."""
        return False

    async def _send_state_command(self, url: str, command: str) -> None:
        """Instruct the API to change the state of the device."""
        # If the user tries to open or close, say, a gateway, throw an exception:
        if not self.state:
            raise RequestError(
                f"Cannot change state of device type: {self.device_type}"
            )

        _LOGGER.debug(f"Sending command {command} for {self.name}")
        await self._api.request(
            method="put",
            returns="response",
            url=url,
        )

    async def update(self) -> None:
        """Get the latest info for this device."""
        await self._api.update_device_info()

    async def wait_for_state(
        self,
        current_state: List,
        new_state: List,
        last_state_update: datetime,
        timeout: int = WAIT_TIMEOUT,
    ) -> bool:
        # First wait until door state is actually updated.
        _LOGGER.debug(f"Waiting until device state has been updated for {self.name}")
        wait_timeout = timeout
        while (
            last_state_update
            == self.device_json["state"].get("last_update", datetime.utcnow())
            and wait_timeout > 0
        ):
            wait_timeout = wait_timeout - 5
            await asyncio.sleep(5)
            try:
                await self._api.update_device_info(for_account=self.account)
            except MyQError:
                # Ignoring
                pass

        # Wait until the state is to what we want it to be
        _LOGGER.debug(f"Waiting until device state for {self.name} is {new_state}")
        wait_timeout = timeout
        while self.state in current_state and wait_timeout > 0:
            wait_timeout = wait_timeout - 5
            await asyncio.sleep(5)
            try:
                await self._api.update_device_info(for_account=self.account)
            except MyQError:
                # Ignoring
                pass

        # Reset self.state ensuring it reflects actual device state. Only do this if state is still what it would
        # have been, this to ensure if something else had updated it to something else we don't override.
        if self._device_state == current_state:
            self._device_state = None

        return self.state in new_state
<|endoftext|>


<commit_before>#!/usr/bin/env python3
## -*- coding: utf-8 -*-
## Copyright (C) 2014 Petri Heinilä, LGPL 2.1

import sys      # http://docs.python.org/py3k/library/sys.html
import argparse # http://docs.python.org/py3k/library/argparse.html
import logging  # http://docs.python.org/py3k/library/logging.html
import datetime
log = logging.getLogger(__name__)
D = log.debug
I = log.info
E = log.error
from pynames.ast_visit import Visitor
from pynames.stats import Stats
import ast
import jinja2

##############################################################################

def extract(file, stats):
  with open(file) as fo:
    code = fo.read()
  tree = ast.parse(code)
  v = Visitor(stats.add_name, file)
  v.visit(tree)

def report(stats):
  print("\n".join(["%s\t%s" % (stats.names[x],x) for x in sorted(stats.names)]))

def report_html(stats, html_file):
  jenv = jinja2.Environment(loader=jinja2.PackageLoader("pynames", "default"))
  tmpl = jenv.get_template("file.html")
  ctx = {
    "page": {
       "title": "pynames",
       "time": datetime.datetime.now().isoformat(),
     },
    "stats": stats
  }
  log.info("generating %s", html_file)
  with open(html_file,"w") as fd:
    fd.write(tmpl.render(ctx))

##############################################################################

ARGS = argparse.ArgumentParser()
ARGS.add_argument("files",
                  metavar="file.py",
                  nargs="+",
                  help="python files to lookup names")
ARGS.add_argument("-d", "--debug",
                  action="store_true",
                  help="set debugging on")
ARGS.add_argument("--html",
                  metavar="FILE",
                  help="produce html file")

def main():
  args = ARGS.parse_args()
  logging.basicConfig()
  if args.debug:
    logging.getLogger().setLevel(logging.DEBUG)
  stats = Stats()
  for file in args.files:
    extract(file, stats)
  if args.html:
    report_html(stats, args.html)
  else:
    report(stats)
  sys.exit(0)

if __name__ == "__main__":
  main()
  <commit_msg>fix on console report<commit_after>#!/usr/bin/env python3
## -*- coding: utf-8 -*-
## Copyright (C) 2014 Petri Heinilä, LGPL 2.1

import sys      # http://docs.python.org/py3k/library/sys.html
import argparse # http://docs.python.org/py3k/library/argparse.html
import logging  # http://docs.python.org/py3k/library/logging.html
import datetime
log = logging.getLogger(__name__)
D = log.debug
I = log.info
E = log.error
from pynames.ast_visit import Visitor
from pynames.stats import Stats
import ast
import jinja2

##############################################################################

def extract(file, stats):
  with open(file) as fo:
    code = fo.read()
  tree = ast.parse(code)
  v = Visitor(stats.add_name, file)
  v.visit(tree)

def report(stats):
  print("\n".join(["%s\t%s" % (stats.names[x].count,x) for x in sorted(stats.names)]))

def report_html(stats, html_file):
  jenv = jinja2.Environment(loader=jinja2.PackageLoader("pynames", "default"))
  tmpl = jenv.get_template("file.html")
  ctx = {
    "page": {
       "title": "pynames",
       "time": datetime.datetime.now().isoformat(),
     },
    "stats": stats
  }
  log.info("generating %s", html_file)
  with open(html_file,"w") as fd:
    fd.write(tmpl.render(ctx))

##############################################################################

ARGS = argparse.ArgumentParser()
ARGS.add_argument("files",
                  metavar="file.py",
                  nargs="+",
                  help="python files to lookup names")
ARGS.add_argument("-d", "--debug",
                  action="store_true",
                  help="set debugging on")
ARGS.add_argument("--html",
                  metavar="FILE",
                  help="produce html file")

def main():
  args = ARGS.parse_args()
  logging.basicConfig()
  if args.debug:
    logging.getLogger().setLevel(logging.DEBUG)
  stats = Stats()
  for file in args.files:
    extract(file, stats)
  if args.html:
    report_html(stats, args.html)
  else:
    report(stats)
  sys.exit(0)

if __name__ == "__main__":
  main()
  <|endoftext|>