<commit_before>##########################################################################
#
#  Copyright (c) 2010, Image Engine Design Inc. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#     * Neither the name of Image Engine Design nor the names of any
#       other contributors to this software may be used to endorse or
#       promote products derived from this software without specific prior
#       written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
#  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
#  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
#  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##########################################################################

import os
from IECore import warning, msg, Msg, error, Writer, CompoundParameter, ClassLoader, SearchPath

## This class manages loading and saving named preset values for Parameterised objects.
# It uses a given SearchPath to locate the presets. The presets must be classes which the
# __call__ method is implemented and the only parameter passed is the Parameterised object to be modified.
class PresetManager :

	## Creates a PresetManager which will load parameter presets using the given classLoader.
	# \param useClassPath If true, then it uses the 'path' member added to the Parameterised objects
	# when loaded by ClassLoader. Otherwise it will only use the name of the Parameterised object.
	def __init__( self, searchPaths, useClassPath = True ) :
		self.__useClassPath = useClassPath
		self.__searchPaths = searchPaths
		self.__classLoader = ClassLoader( searchPaths )

	## Returns a copy of the SearchPath object given on the constructor.
	def searchPaths( self ):
		return SearchPath( ":".join( self.__searchPaths.paths ), ":" )

	## Refreshes the internal cache for the preset loader.
	def refresh( self ) :
		self.__classLoader.refresh()

	## Returns the preset names available for a given parameterised object.
	def presets( self, parameterised ):
		presets = self.__classLoader.classNames( self.__presetName( parameterised, "*" ) )
		return [ p.split("/")[-1] for p in presets ]

	## Converts from user preset name to class name
	def __presetName( self, parameterised, name ):
		if self.__useClassPath :
			return "%s/%s" % ( parameterised.path, name )
		else:
			return "%s/%s" % ( parameterised.typeName(), name )

	# This function uses the class name and version from the given parameterised object
	# to find the correspondent presetName and update the parameterised parameters that match.
	# It assumes the given parameterised object was loaded by ClassLoader and so, have the attributes 'path' and 'version'.
	def loadPreset( self, parameterised, presetName ) :

		preset = self.__classLoader.load( self.__presetName( parameterised, presetName ) )()
		preset( parameterised )

	# Saves the parameter values for the given Parameterised object only for the given Parameter list as a named preset.
	# The path should point to a valid location ( used by the internal SearchPath )
	def savePreset( self, parameterised, parameterList, path, presetName ) :

		# in order to match ClassLoader convention, the final path contains the presetVersion (constant) and has .py extension.
		presetPath = os.path.join( path, self.__presetName( parameterised, presetName ) )
		presetFile = os.path.join( presetPath, presetName + "-1.py" )

		tmpParameterList = list( parameterList )
		parameterPaths = {}

		def saveParam( param, curParamPath ):
			myPath = list(curParamPath)
			myPath.append( param.name )
			# ignore root CompoundParameter name
			del( myPath[0] )
			parameterPaths[ param ] = myPath
			Writer.create( param.getValue(), os.path.join( presetPath, presetName + "." + ".".join( myPath ) + ".cob" ) )()

		def findParamIndex( param, paramList ):
			index = 0
			for p in paramList:
				if p.isSame( param ):
					return index
				index += 1
			return -1

		def recursiveSave( compoundParam, curParamPath ):

			paramIndex = findParamIndex( compoundParam, tmpParameterList )
			if paramIndex != -1:
				saveParam( compoundParam, curParamPath )
				del( tmpParameterList[ paramIndex ] )
				return

			newParamPath = list( curParamPath )
			newParamPath.append( compoundParam.name )

			for param in compoundParam.values():

				if isinstance( param, CompoundParameter ):
					recursiveSave( param, newParamPath )
				else:
					paramIndex = findParamIndex( param, tmpParameterList )
					if param in tmpParameterList :
						saveParam( param, newParamPath )
						del( tmpParameterList[ paramIndex ] )

		# make sure the preset dir exists
		if not os.path.exists( presetPath ):
			os.makedirs( presetPath )

		# save COB files
		recursiveSave( parameterised.parameters(), [] )

		# save the preset script.
		presetFile = open( presetFile, "w" )
		presetFile.write( """
import IECore
class %s:
	def __call__( self, parameterised ):
""" % presetName
		)
		parameterPathNames = parameterPaths.keys()
		for p in parameterList :
			# find same parameter in dictionary using isSame method...
			paramPath = parameterPaths[ parameterPathNames[findParamIndex(p, parameterPathNames)] ]
			presetFile.write( "\t\tIECore.PresetManager.applyParameterValue( parameterised, %s, IECore.Reader.create( '%s' )() )\n" %
				( paramPath, os.path.join( presetPath, presetName + "." + ".".join( paramPath ) + ".cob" ) ) )

		presetFile.close()

	# Utility function that is used by standard preset scripts written by this same very PresetManager.
	@staticmethod
	def applyParameterValue( parameterised, parameterPath, parameterValue ):

		param = parameterised.parameters()
		for paramName in parameterPath :
			if not paramName in param:
				warning( "Ignoring parameter %s from preset." % ( ".".join( parameterPath ) ) )
				continue
			param = param[ paramName ]

		try:
			param.setValue( parameterValue )
		except Exception, e:
			warning( "Could not set preset value for parameter %s: %s" % ( ".".join( parameterPath ), str(e) ) )

	__defaultManagers = {}

	## Returns a PresetManager configured to load from the paths defined by the
	# specified environment variable. The same object is returned each time,
	# allowing one manager to be shared by many callers.
	# It uses ClassLoader.defaultLoader() method.
	@classmethod
	def defaultManager( cls, envVar ) :

		mgr = cls.__defaultManagers.get( envVar, None )
		if mgr :
			return mgr

		sp = ""
		if envVar in os.environ :
			sp = os.environ[envVar]
		else :
			msg( Msg.Level.Warning, "PresetManager.defaultLoader", "Environment variable %s not set." % envVar )

		manager = cls( SearchPath( os.path.expandvars( sp ), ":" ) )
		cls.__defaultManagers[envVar] = manager

		return manager
<commit_msg>Making sure that we don't pass a unicode string to the binding, which fails.<commit_after>##########################################################################
#
#  Copyright (c) 2010, Image Engine Design Inc. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#     * Neither the name of Image Engine Design nor the names of any
#       other contributors to this software may be used to endorse or
#       promote products derived from this software without specific prior
#       written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
#  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
#  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
#  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##########################################################################

import os
from IECore import warning, msg, Msg, error, Writer, CompoundParameter, ClassLoader, SearchPath

## This class manages loading and saving named preset values for Parameterised objects.
# It uses a given SearchPath to locate the presets. The presets must be classes which the
# __call__ method is implemented and the only parameter passed is the Parameterised object to be modified.
class PresetManager :

	## Creates a PresetManager which will load parameter presets using the given classLoader.
	# \param useClassPath If true, then it uses the 'path' member added to the Parameterised objects
	# when loaded by ClassLoader. Otherwise it will only use the name of the Parameterised object.
	def __init__( self, searchPaths, useClassPath = True ) :
		self.__useClassPath = useClassPath
		self.__searchPaths = searchPaths
		self.__classLoader = ClassLoader( searchPaths )

	## Returns a copy of the SearchPath object given on the constructor.
	def searchPaths( self ):
		return SearchPath( ":".join( self.__searchPaths.paths ), ":" )

	## Refreshes the internal cache for the preset loader.
	def refresh( self ) :
		self.__classLoader.refresh()

	## Returns the preset names available for a given parameterised object.
	def presets( self, parameterised ):
		presets = self.__classLoader.classNames( self.__presetName( parameterised, "*" ) )
		return [ p.split("/")[-1] for p in presets ]

	## Converts from user preset name to class name
	def __presetName( self, parameterised, name ):
		if self.__useClassPath :
			return "%s/%s" % ( parameterised.path, name )
		else:
			return "%s/%s" % ( parameterised.typeName(), name )

	# This function uses the class name and version from the given parameterised object
	# to find the correspondent presetName and update the parameterised parameters that match.
	# It assumes the given parameterised object was loaded by ClassLoader and so, have the attributes 'path' and 'version'.
	def loadPreset( self, parameterised, presetName ) :
		preset = self.__classLoader.load( self.__presetName( parameterised, str(presetName) ) )()
		preset( parameterised )

	# Saves the parameter values for the given Parameterised object only for the given Parameter list as a named preset.
	# The path should point to a valid location ( used by the internal SearchPath )
	def savePreset( self, parameterised, parameterList, path, presetName ) :

		# in order to match ClassLoader convention, the final path contains the presetVersion (constant) and has .py extension.
		presetPath = os.path.join( path, self.__presetName( parameterised, presetName ) )
		presetFile = os.path.join( presetPath, presetName + "-1.py" )

		tmpParameterList = list( parameterList )
		parameterPaths = {}

		def saveParam( param, curParamPath ):
			myPath = list(curParamPath)
			myPath.append( param.name )
			# ignore root CompoundParameter name
			del( myPath[0] )
			parameterPaths[ param ] = myPath
			Writer.create( param.getValue(), os.path.join( presetPath, presetName + "." + ".".join( myPath ) + ".cob" ) )()

		def findParamIndex( param, paramList ):
			index = 0
			for p in paramList:
				if p.isSame( param ):
					return index
				index += 1
			return -1

		def recursiveSave( compoundParam, curParamPath ):

			paramIndex = findParamIndex( compoundParam, tmpParameterList )
			if paramIndex != -1:
				saveParam( compoundParam, curParamPath )
				del( tmpParameterList[ paramIndex ] )
				return

			newParamPath = list( curParamPath )
			newParamPath.append( compoundParam.name )

			for param in compoundParam.values():

				if isinstance( param, CompoundParameter ):
					recursiveSave( param, newParamPath )
				else:
					paramIndex = findParamIndex( param, tmpParameterList )
					if param in tmpParameterList :
						saveParam( param, newParamPath )
						del( tmpParameterList[ paramIndex ] )

		# make sure the preset dir exists
		if not os.path.exists( presetPath ):
			os.makedirs( presetPath )

		# save COB files
		recursiveSave( parameterised.parameters(), [] )

		# save the preset script.
		presetFile = open( presetFile, "w" )
		presetFile.write( """
import IECore
class %s:
	def __call__( self, parameterised ):
""" % presetName
		)
		parameterPathNames = parameterPaths.keys()
		for p in parameterList :
			# find same parameter in dictionary using isSame method...
			paramPath = parameterPaths[ parameterPathNames[findParamIndex(p, parameterPathNames)] ]
			presetFile.write( "\t\tIECore.PresetManager.applyParameterValue( parameterised, %s, IECore.Reader.create( '%s' )() )\n" %
				( paramPath, os.path.join( presetPath, presetName + "." + ".".join( paramPath ) + ".cob" ) ) )

		presetFile.close()

	# Utility function that is used by standard preset scripts written by this same very PresetManager.
	@staticmethod
	def applyParameterValue( parameterised, parameterPath, parameterValue ):

		param = parameterised.parameters()
		for paramName in parameterPath :
			if not paramName in param:
				warning( "Ignoring parameter %s from preset." % ( ".".join( parameterPath ) ) )
				continue
			param = param[ paramName ]

		try:
			param.setValue( parameterValue )
		except Exception, e:
			warning( "Could not set preset value for parameter %s: %s" % ( ".".join( parameterPath ), str(e) ) )

	__defaultManagers = {}

	## Returns a PresetManager configured to load from the paths defined by the
	# specified environment variable. The same object is returned each time,
	# allowing one manager to be shared by many callers.
	# It uses ClassLoader.defaultLoader() method.
	@classmethod
	def defaultManager( cls, envVar ) :

		mgr = cls.__defaultManagers.get( envVar, None )
		if mgr :
			return mgr

		sp = ""
		if envVar in os.environ :
			sp = os.environ[envVar]
		else :
			msg( Msg.Level.Warning, "PresetManager.defaultLoader", "Environment variable %s not set." % envVar )

		manager = cls( SearchPath( os.path.expandvars( sp ), ":" ) )
		cls.__defaultManagers[envVar] = manager

		return manager
<|endoftext|>


<commit_before>from typing import BinaryIO, List, Union

import numpy as np
import numpy.ma as ma
from scipy.interpolate import interp1d

from ..pose_body import PoseBody, POINTS_DIMS
from ..pose_header import PoseHeader
from ..utils.reader import BufferReader, ConstStructs

# import numpy as np
# np.seterr(all='raise')

class NumPyPoseBody(PoseBody):
    tensor_reader = 'unpack_numpy'

    def __init__(self, fps: int, data: Union[ma.MaskedArray, np.ndarray], confidence: np.ndarray):
        if isinstance(data, np.ndarray):  # If array is not masked
            mask = confidence == 0  # 0 means no-mask, 1 means with-mask
            stacked_mask = np.stack([mask, mask], axis=3)
            data = ma.masked_array(data, mask=stacked_mask)

        super().__init__(fps, data, confidence)

    @classmethod
    def read_v0_0(cls, header: PoseHeader, reader: BufferReader):
        fps, _frames = reader.unpack(ConstStructs.double_ushort)

        _dims = max([len(c.format) for c in header.components]) - 1
        _points = sum([len(c.points) for c in header.components])

        frames_d = []
        frames_c = []
        for _ in range(_frames):
            _people = reader.unpack(ConstStructs.ushort)
            people_d = []
            people_c = []
            for pid in range(_people):
                reader.advance(ConstStructs.short)  # Skip Person ID
                person_d = []
                person_c = []
                for component in header.components:
                    points = np.array(
                        reader.unpack_numpy(ConstStructs.float, (len(component.points), len(component.format))))
                    dimensions, confidence = np.split(points, [-1], axis=1)
                    boolean_confidence = np.where(confidence > 0, 0, 1)  # To create the mask
                    mask = np.column_stack(tuple([boolean_confidence] * (len(component.format) - 1)))

                    person_d.append(ma.masked_array(dimensions, mask=mask))
                    person_c.append(np.squeeze(confidence, axis=-1))

                if pid == 0:
                    people_d.append(ma.concatenate(person_d))
                    people_c.append(np.concatenate(person_c))

            # In case no person, should all be zeros
            if len(people_d) == 0:
                people_d.append(np.zeros((_points, _dims)))
                people_c.append(np.zeros(_points))

            frames_d.append(ma.stack(people_d))
            frames_c.append(np.stack(people_c))

        return cls(fps, ma.stack(frames_d), ma.stack(frames_c))

    def write(self, version: float, buffer: BinaryIO):
        _frames, _people, _points, _dims = self.data.shape
        _frames = _frames if _frames < 65535 else 0  # TODO change from short to int
        buffer.write(ConstStructs.triple_ushort.pack(self.fps, _frames, _people))

        buffer.write(np.array(self.data.data, dtype=np.float32).tobytes())
        buffer.write(np.array(self.confidence, dtype=np.float32).tobytes())

    @property
    def mask(self):
        return self.data.mask

    def torch(self):
        import torch
        from ..torch.pose_body import TorchPoseBody

        torch_confidence = torch.from_numpy(self.confidence)
        torch_data = torch.from_numpy(self.data.data)
        return TorchPoseBody(self.fps, torch_data, torch_confidence)

    def tensorflow(self):
        import tensorflow
        from ..tensorflow.pose_body import TensorflowPoseBody

        tf_confidence = tensorflow.constant(self.confidence)
        tf_data = tensorflow.constant(self.data.data)
        return TensorflowPoseBody(self.fps, tf_data, tf_confidence)

    def zero_filled(self):
        self.data = self.data.filled(0)
        return self

    def matmul(self, matrix: np.ndarray):
        data = ma.dot(self.data, matrix)
        return NumPyPoseBody(self.fps, data, self.confidence)

    def points_perspective(self):
        return ma.transpose(self.data, axes=POINTS_DIMS)

    def get_points(self, indexes: List[int]):
        data = ma.transpose(self.data, axes=POINTS_DIMS)
        new_data = ma.transpose(data[indexes], axes=POINTS_DIMS)

        confidence_reshape = (2, 1, 0)
        confidence = np.transpose(self.confidence, axes=confidence_reshape)
        new_confidence = np.transpose(confidence[indexes], axes=confidence_reshape)

        return NumPyPoseBody(self.fps, new_data, new_confidence)

    def bbox(self, header: PoseHeader):
        data = ma.transpose(self.data, axes=POINTS_DIMS)

        # Split data by components, `ma` doesn't support ".split"
        components = []
        idx = 0
        for component in header.components:
            components.append(data[list(range(idx, idx + len(component.points)))])
            idx += len(component.points)

        boxes = [ma.stack([ma.min(c, axis=0), ma.max(c, axis=0)]) for c in components]
        boxes_cat = ma.concatenate(boxes)
        if type(boxes_cat.mask) == np.bool_:  # Sometimes, it doesn't concatenate the mask...
            boxes_mask = ma.concatenate([b.mask for b in boxes])
            boxes_cat = ma.array(boxes_cat, mask=boxes_mask)

        new_data = ma.transpose(boxes_cat, axes=POINTS_DIMS)

        confidence_mask = np.split(new_data.mask, [-1], axis=3)[0]
        confidence_mask = np.squeeze(confidence_mask, axis=-1)
        confidence = np.where(confidence_mask == True, 0, 1)

        return NumPyPoseBody(self.fps, new_data, confidence)

    def interpolate(self, new_fps: int = None, kind='cubic'):
        if new_fps is None:
            new_fps = self.fps

        _frames = self.data.shape[0]
        if _frames == 1:
            raise ValueError("Can't interpolate single frame")

        _new_frames = round(_frames * new_fps / self.fps)
        steps = np.linspace(0, 1, _frames)
        new_steps = np.linspace(0, 1, _new_frames)

        transposed = self.points_perspective()  # (points, people, frames, dims)
        masked_confidence = ma.array(self.confidence, mask=self.confidence == 0)
        confidence = ma.expand_dims(masked_confidence.transpose(), axis=3)  # (points, people, frames, 1)
        points = ma.concatenate([transposed, confidence], axis=3)

        new_people = []
        for people in points:
            new_frames = []
            for frames in people:
                mask = frames.transpose()[0].mask

                partial_steps = ma.array(steps, mask=mask).compressed()

                if partial_steps.shape[0] == 0:  # No data for this point
                    new_frames.append(np.zeros((_new_frames, frames.shape[1])))
                else:
                    partial_frames = frames.compressed().reshape(partial_steps.shape[0], frames.shape[1])

                    if len(partial_steps) == 1:
                        f = lambda l: partial_frames
                    else:
                        this_kind = kind if len(partial_steps) > 3 \
                            else "quadratic" if len(partial_steps) > 2 and kind == "cubic" \
                            else "linear"  # Can't do something fancy for 2 points
                        f = interp1d(partial_steps, partial_frames, axis=0, kind=this_kind)

                    first_step = partial_steps[0]
                    last_step = partial_steps[-1]
                    if first_step == 0 and last_step == 1:
                        new_frames.append(f(new_steps))
                    else:
                        first_step_where = np.argwhere(new_steps >= first_step)
                        first_step_index = first_step_where[0][0] if len(first_step_where) > 0 else 0

                        last_step_where = np.argwhere(new_steps > last_step)
                        last_step_index = last_step_where[0][0] if len(last_step_where) > 0 else len(new_steps)

                        if first_step_index == last_step_index:
                            new_frames.append(np.zeros((len(new_steps), frames.shape[1])))
                        else:
                            frame_data = f(new_steps[first_step_index:last_step_index])
                            new_frames.append(np.concatenate([
                                np.zeros((first_step_index, frames.shape[1])),
                                np.array(frame_data),
                                np.zeros((len(new_steps) - last_step_index, frames.shape[1]))
                            ]))
            new_people.append(np.stack(new_frames, axis=0))

        new_data = np.stack(new_people, axis=0).transpose([2, 1, 0, 3])
        dimensions, confidence = np.split(new_data, [-1], axis=3)
        confidence = np.squeeze(confidence, axis=3)

        mask = confidence == 0
        stacked_confidence = np.stack([mask, mask], axis=3)
        masked_data = ma.masked_array(dimensions, mask=stacked_confidence)

        return NumPyPoseBody(fps=new_fps, data=masked_data, confidence=confidence)

    def flatten(self):
        shape = self.data.shape
        data = self.data.data.reshape(-1, shape[-1])  # Not masked data
        confidence = self.confidence.flatten()
        indexes = list(np.ndindex(shape[:-1]))
        flat = np.c_[indexes, confidence, data]
        # Filter data from flat
        flat = flat[confidence != 0]
        # Scale the first axis by fps
        scalar = np.ones(len(shape) + shape[-1])
        scalar[0] = 1 / self.fps
        return flat * scalar
<commit_msg>fix(pose_body): stack body correctly<commit_after>from typing import BinaryIO, List, Union

import numpy as np
import numpy.ma as ma
from scipy.interpolate import interp1d

from ..pose_body import PoseBody, POINTS_DIMS
from ..pose_header import PoseHeader
from ..utils.reader import BufferReader, ConstStructs

# import numpy as np
# np.seterr(all='raise')

class NumPyPoseBody(PoseBody):
    tensor_reader = 'unpack_numpy'

    def __init__(self, fps: int, data: Union[ma.MaskedArray, np.ndarray], confidence: np.ndarray):
        if isinstance(data, np.ndarray):  # If array is not masked
            mask = confidence == 0  # 0 means no-mask, 1 means with-mask
            stacked_mask = np.stack([mask] * data.shape[-1], axis=3)
            data = ma.masked_array(data, mask=stacked_mask)

        super().__init__(fps, data, confidence)

    @classmethod
    def read_v0_0(cls, header: PoseHeader, reader: BufferReader):
        fps, _frames = reader.unpack(ConstStructs.double_ushort)

        _dims = max([len(c.format) for c in header.components]) - 1
        _points = sum([len(c.points) for c in header.components])

        frames_d = []
        frames_c = []
        for _ in range(_frames):
            _people = reader.unpack(ConstStructs.ushort)
            people_d = []
            people_c = []
            for pid in range(_people):
                reader.advance(ConstStructs.short)  # Skip Person ID
                person_d = []
                person_c = []
                for component in header.components:
                    points = np.array(
                        reader.unpack_numpy(ConstStructs.float, (len(component.points), len(component.format))))
                    dimensions, confidence = np.split(points, [-1], axis=1)
                    boolean_confidence = np.where(confidence > 0, 0, 1)  # To create the mask
                    mask = np.column_stack(tuple([boolean_confidence] * (len(component.format) - 1)))

                    person_d.append(ma.masked_array(dimensions, mask=mask))
                    person_c.append(np.squeeze(confidence, axis=-1))

                if pid == 0:
                    people_d.append(ma.concatenate(person_d))
                    people_c.append(np.concatenate(person_c))

            # In case no person, should all be zeros
            if len(people_d) == 0:
                people_d.append(np.zeros((_points, _dims)))
                people_c.append(np.zeros(_points))

            frames_d.append(ma.stack(people_d))
            frames_c.append(np.stack(people_c))

        return cls(fps, ma.stack(frames_d), ma.stack(frames_c))

    def write(self, version: float, buffer: BinaryIO):
        _frames, _people, _points, _dims = self.data.shape
        _frames = _frames if _frames < 65535 else 0  # TODO change from short to int
        buffer.write(ConstStructs.triple_ushort.pack(self.fps, _frames, _people))

        buffer.write(np.array(self.data.data, dtype=np.float32).tobytes())
        buffer.write(np.array(self.confidence, dtype=np.float32).tobytes())

    @property
    def mask(self):
        return self.data.mask

    def torch(self):
        import torch
        from ..torch.pose_body import TorchPoseBody

        torch_confidence = torch.from_numpy(self.confidence)
        torch_data = torch.from_numpy(self.data.data)
        return TorchPoseBody(self.fps, torch_data, torch_confidence)

    def tensorflow(self):
        import tensorflow
        from ..tensorflow.pose_body import TensorflowPoseBody

        tf_confidence = tensorflow.constant(self.confidence)
        tf_data = tensorflow.constant(self.data.data)
        return TensorflowPoseBody(self.fps, tf_data, tf_confidence)

    def zero_filled(self):
        self.data = self.data.filled(0)
        return self

    def matmul(self, matrix: np.ndarray):
        data = ma.dot(self.data, matrix)
        return NumPyPoseBody(self.fps, data, self.confidence)

    def flip(self, axis=0):
        vec = np.ones(self.data.shape[-1])
        vec[axis] = -1

        data = self.data * vec
        return NumPyPoseBody(self.fps, data, self.confidence)

    def points_perspective(self):
        return ma.transpose(self.data, axes=POINTS_DIMS)

    def get_points(self, indexes: List[int]):
        data = ma.transpose(self.data, axes=POINTS_DIMS)
        new_data = ma.transpose(data[indexes], axes=POINTS_DIMS)

        confidence_reshape = (2, 1, 0)
        confidence = np.transpose(self.confidence, axes=confidence_reshape)
        new_confidence = np.transpose(confidence[indexes], axes=confidence_reshape)

        return NumPyPoseBody(self.fps, new_data, new_confidence)

    def bbox(self, header: PoseHeader):
        data = ma.transpose(self.data, axes=POINTS_DIMS)

        # Split data by components, `ma` doesn't support ".split"
        components = []
        idx = 0
        for component in header.components:
            components.append(data[list(range(idx, idx + len(component.points)))])
            idx += len(component.points)

        boxes = [ma.stack([ma.min(c, axis=0), ma.max(c, axis=0)]) for c in components]
        boxes_cat = ma.concatenate(boxes)
        if type(boxes_cat.mask) == np.bool_:  # Sometimes, it doesn't concatenate the mask...
            boxes_mask = ma.concatenate([b.mask for b in boxes])
            boxes_cat = ma.array(boxes_cat, mask=boxes_mask)

        new_data = ma.transpose(boxes_cat, axes=POINTS_DIMS)

        confidence_mask = np.split(new_data.mask, [-1], axis=3)[0]
        confidence_mask = np.squeeze(confidence_mask, axis=-1)
        confidence = np.where(confidence_mask == True, 0, 1)

        return NumPyPoseBody(self.fps, new_data, confidence)

    def interpolate(self, new_fps: int = None, kind='cubic'):
        if new_fps is None:
            new_fps = self.fps

        _frames = self.data.shape[0]
        if _frames == 1:
            raise ValueError("Can't interpolate single frame")

        _new_frames = round(_frames * new_fps / self.fps)
        steps = np.linspace(0, 1, _frames)
        new_steps = np.linspace(0, 1, _new_frames)

        transposed = self.points_perspective()  # (points, people, frames, dims)
        masked_confidence = ma.array(self.confidence, mask=self.confidence == 0)
        confidence = ma.expand_dims(masked_confidence.transpose(), axis=3)  # (points, people, frames, 1)
        points = ma.concatenate([transposed, confidence], axis=3)

        new_people = []
        for people in points:
            new_frames = []
            for frames in people:
                mask = frames.transpose()[0].mask

                partial_steps = ma.array(steps, mask=mask).compressed()

                if partial_steps.shape[0] == 0:  # No data for this point
                    new_frames.append(np.zeros((_new_frames, frames.shape[1])))
                else:
                    partial_frames = frames.compressed().reshape(partial_steps.shape[0], frames.shape[1])

                    if len(partial_steps) == 1:
                        f = lambda l: partial_frames
                    else:
                        this_kind = kind if len(partial_steps) > 3 \
                            else "quadratic" if len(partial_steps) > 2 and kind == "cubic" \
                            else "linear"  # Can't do something fancy for 2 points
                        f = interp1d(partial_steps, partial_frames, axis=0, kind=this_kind)

                    first_step = partial_steps[0]
                    last_step = partial_steps[-1]
                    if first_step == 0 and last_step == 1:
                        new_frames.append(f(new_steps))
                    else:
                        first_step_where = np.argwhere(new_steps >= first_step)
                        first_step_index = first_step_where[0][0] if len(first_step_where) > 0 else 0

                        last_step_where = np.argwhere(new_steps > last_step)
                        last_step_index = last_step_where[0][0] if len(last_step_where) > 0 else len(new_steps)

                        if first_step_index == last_step_index:
                            new_frames.append(np.zeros((len(new_steps), frames.shape[1])))
                        else:
                            frame_data = f(new_steps[first_step_index:last_step_index])
                            new_frames.append(np.concatenate([
                                np.zeros((first_step_index, frames.shape[1])),
                                np.array(frame_data),
                                np.zeros((len(new_steps) - last_step_index, frames.shape[1]))
                            ]))
            new_people.append(np.stack(new_frames, axis=0))

        new_data = np.stack(new_people, axis=0).transpose([2, 1, 0, 3])
        dimensions, confidence = np.split(new_data, [-1], axis=3)
        confidence = np.squeeze(confidence, axis=3)

        return NumPyPoseBody(fps=new_fps, data=dimensions, confidence=confidence)

    def flatten(self):
        shape = self.data.shape
        data = self.data.data.reshape(-1, shape[-1])  # Not masked data
        confidence = self.confidence.flatten()
        indexes = list(np.ndindex(shape[:-1]))
        flat = np.c_[indexes, confidence, data]
        # Filter data from flat
        flat = flat[confidence != 0]
        # Scale the first axis by fps
        scalar = np.ones(len(shape) + shape[-1])
        scalar[0] = 1 / self.fps
        return flat * scalar
<|endoftext|>


<commit_before>##
# copyright 2009, James William Pye
# http://python.projects.postgresql.org
##
"""
itertools extensions
"""
import collections
from itertools import cycle

def interlace(*iters) -> collections.Iterable:
	"""
	interlace(i1, i2, ..., in) -> (
		i1-0, i2-0, ..., in-0,
		i1-1, i2-1, ..., in-1,
		.
		.
		.
		i1-n, i2-n, ..., in-n,
	)
	"""
	iters = [iter(x) for x in iters]
	for i in cycle(range(len(iters))):
		yield next(iters[i])
<commit_msg>Rewrite interlace using map(next,cycle())<commit_after>##
# copyright 2009, James William Pye
# http://python.projects.postgresql.org
##
"""
itertools extensions
"""
import collections
from itertools import cycle

def interlace(*iters) -> collections.Iterable:
	"""
	interlace(i1, i2, ..., in) -> (
		i1-0, i2-0, ..., in-0,
		i1-1, i2-1, ..., in-1,
		.
		.
		.
		i1-n, i2-n, ..., in-n,
	)
	"""
	return map(next, cycle([iter(x) for x in iters]))
<|endoftext|>


<commit_before>#!/usr/bin/evn python
# encoding: utf-8

from pprof.project import (ProjectFactory, log_with, log)
from pprof.settings import config
from group import PprofGroup

from os import path
from plumbum import FG, local
from plumbum.cmd import ln

class Ccrypt(PprofGroup):

    """ ccrypt benchmark """

    check_f = "check"

    class Factory:

        def create(self, exp):
            return Ccrypt(exp, "ccrypt", "encryption")
    ProjectFactory.addFactory("Ccrypt", Factory())

    def prepare(self):
        super(Ccrypt, self).prepare()
        check_f = path.join(self.testdir, self.check_f)
        ln("-s", check_f, path.join(self.builddir, self.check_f))

    def clean(self):
        check_f = path.join(self.builddir, self.check_f)
        self.products.add(check_f)

        super(Ccrypt, self).clean()

    src_dir = "ccrypt-1.10"
    src_file = "ccrypt-1.10.tar.gz"
    src_uri = "http://ccrypt.sourceforge.net/download/ccrypt-1.10.tar.gz"

    def download(self):
        from pprof.utils.downloader import Wget
        from plumbum.cmd import tar, cp

        ccrypt_dir = path.join(self.builddir, self.src_dir)
        with local.cwd(self.builddir):
            Wget(self.src_uri, self.src_file)
            tar('xfz', path.join(self.builddir, self.src_file))

    def configure(self):
        from pprof.utils.compiler import lt_clang, lt_clang_cxx
        ccrypt_dir = path.join(self.builddir, self.src_dir)

        with local.cwd(ccrypt_dir):
            configure = local[path.join(ccrypt_dir, "configure")]

            with local.env(CC=str(lt_clang(self.cflags)),
                           CXX=str(lt_clang_cxx(self.cflags)),
                           LDFLAGS=" ".join(self.ldflags)):
                configure & FG

    def build(self):
        from plumbum.cmd import make, mv, rm

        ccrypt_dir = path.join(self.builddir, self.src_dir)
        with local.cwd(ccrypt_dir):
            sh_file = path.join("src", "ccrypt")
            rm("-f", sh_file)
            make & FG
            mv(sh_file, self.bin_f)
        self.run_f = self.bin_f

    def run_tests(self, experiment):
        from plumbum.cmd import make, chmod

        exp = experiment(self.run_f)

        ccrypt_dir = path.join(self.builddir, self.src_dir)
        with local.cwd(ccrypt_dir):
            sh_file = path.join("src", self.name)
            with open(sh_file, 'w') as ccrypt:
                ccrypt.write("#!/usr/bin/env bash\n")
                ccrypt.write("exec " + str(exp["\"$@\""]))
            chmod("+x", sh_file)
            make["check"] & FG
<commit_msg>ccrypt: use wrap_tool<commit_after>#!/usr/bin/evn python
# encoding: utf-8

from pprof.project import (ProjectFactory, log_with, log)
from pprof.settings import config
from group import PprofGroup

from os import path
from plumbum import FG, local
from plumbum.cmd import ln

class Ccrypt(PprofGroup):

    """ ccrypt benchmark """

    check_f = "check"

    class Factory:

        def create(self, exp):
            return Ccrypt(exp, "ccrypt", "encryption")
    ProjectFactory.addFactory("Ccrypt", Factory())

    def prepare(self):
        super(Ccrypt, self).prepare()
        check_f = path.join(self.testdir, self.check_f)
        ln("-s", check_f, path.join(self.builddir, self.check_f))

    def clean(self):
        check_f = path.join(self.builddir, self.check_f)
        self.products.add(check_f)

        super(Ccrypt, self).clean()

    src_dir = "ccrypt-1.10"
    src_file = "ccrypt-1.10.tar.gz"
    src_uri = "http://ccrypt.sourceforge.net/download/ccrypt-1.10.tar.gz"

    def download(self):
        from pprof.utils.downloader import Wget
        from plumbum.cmd import tar, cp

        ccrypt_dir = path.join(self.builddir, self.src_dir)
        with local.cwd(self.builddir):
            Wget(self.src_uri, self.src_file)
            tar('xfz', path.join(self.builddir, self.src_file))

    def configure(self):
        from pprof.utils.compiler import lt_clang, lt_clang_cxx
        ccrypt_dir = path.join(self.builddir, self.src_dir)

        with local.cwd(ccrypt_dir):
            configure = local[path.join(ccrypt_dir, "configure")]
            with local.env(CC=str(lt_clang(self.cflags)),
                           CXX=str(lt_clang_cxx(self.cflags)),
                           LDFLAGS=" ".join(self.ldflags)):
                configure()

    def build(self):
        from plumbum.cmd import make, mv, rm

        ccrypt_dir = path.join(self.builddir, self.src_dir)
        with local.cwd(ccrypt_dir):
            make()

    def run_tests(self, experiment):
        from plumbum.cmd import make, chmod
        from pprof.project import wrap_tool

        exp = experiment(self.run_f)

        ccrypt_dir = path.join(self.builddir, self.src_dir)
        wrap_tool(path.join(ccrypt_dir, "src", self.name), experiment)

        with local.cwd(ccrypt_dir):
            make("check")
<|endoftext|>


<commit_before>import re
from datetime import datetime, date, time as dt_time
from urllib.parse import parse_qsl, urlparse

import pytz
from dateutil.parser import parse as parse_date

from schematics.exceptions import (
    ValidationError as SchematicsValidationError,
    DataError as SchematicsDataError,
    ConversionError,
)
from schematics.models import Model
from schematics.transforms import Role
from schematics.types import (
    StringType,
    DateTimeType as SchematicsDateTimeType,
    IntType,
    FloatType,
    URLType,
    GeoPointType,
    BooleanType,
    DateType as SchematicsDateType,
)
from schematics.types.compound import ListType as SchematicsListType, ModelType, DictType, PolyModelType
from schematics.types.serializable import serializable

class DateTimeType(SchematicsDateTimeType):
    def to_native(self, value, context=None):
        if isinstance(value, datetime):
            return value
        return parse_date(value)

class DateType(SchematicsDateType):
    def to_native(self, value, context=None):
        if isinstance(value, date):
            return value
        return parse_date(value).date()

class DateOrDateTimeType(SchematicsDateTimeType):
    def validate_tz(self, value, context=None):
        if isinstance(value, datetime):
            return super().validate_tz(value, context)

    def to_primitive(self, value, context=None):
        if isinstance(value, datetime):
            return super().to_primitive(value, context)
        else:
            return value.strftime("%Y-%m-%d")

    def to_native(self, value, context=None):
        if isinstance(value, str):
            try:
                return datetime.strptime(value, "%Y-%m-%d").date()
            except ValueError:
                return parse_date(value)
        if isinstance(value, datetime):
            return value
        return value

class StringModelType(ModelType):
    def _convert(self, value, context):
        if isinstance(value, str):
            return self.model_class._convert(value, context=context)
        else:
            return super(StringModelType, self)._convert(value, context=context)

class StringModel(Model):
    @property
    def import_format(self):  # pragma: nocover
        raise NotImplementedError()

    @property
    def export_format(self):  # pragma: nocover
        raise NotImplementedError()

    def convert(self, raw_data, **kw):
        if isinstance(raw_data, str):
            try:
                raw_data = re.match(self.import_format, str(raw_data)).groupdict()
            except AttributeError:
                raise SchematicsValidationError(f"'{raw_data}' is not a valid format for {self.__class__.__name__}")
        return super(StringModel, self).convert(raw_data, **kw)

    def export(self, *args, **kwargs):
        return self.export_format.format(**super(StringModel, self).export(*args, **kwargs))

class StringListType(SchematicsListType):
    def __init__(self, field, separator=" ", **kwargs):
        self.separator = separator
        super(StringListType, self).__init__(field, **kwargs)

    def _coerce(self, value):
        if isinstance(value, str):
            return value.split(self.separator)
        elif isinstance(value, dict):
            return [value]
        else:
            return super(StringListType, self)._coerce(value)

    def _export(self, *args, **kwargs):
        data = super(StringListType, self)._export(*args, **kwargs)
        return self.separator.join(map(str, data))

class ListType(SchematicsListType):
    def _coerce(self, value):
        if not isinstance(value, (list, tuple)):
            return [value]
        else:
            return super(ListType, self)._coerce(value)

class GeoJSONPointType(GeoPointType):
    def _normalize(self, value):
        return tuple(reversed(super(GeoJSONPointType, self)._normalize(value)))

class Area(StringModel):

    import_format = r"(?P<radius>(\d+(\.\d+)?)(km|mi)|(\d+)(m|ft))@(?P<latitude>[\-\+]?\d+(\.\d+)?),(?P<longitude>[\-\+]?\d+(\.\d+)?)"
    export_format = "{radius}@{latitude},{longitude}"

    radius = StringType(regex=r"((\d+(\.\d+)?)+(km|mi)|(\d+)+(m|ft))", required=True)
    latitude = FloatType(required=True)
    longitude = FloatType(required=True)

class Location(StringModel):

    import_format = r"@(?P<latitude>-?\d+(\.\d+)?),(?P<longitude>-?\d+(\.\d+)?)"
    export_format = "@{latitude},{longitude}"

    latitude = FloatType(required=True)
    longitude = FloatType(required=True)

class Place(Model):
    class Options:
        serialize_when_none = False

    scope = ListType(StringType)
    exact = ListType(StringType)

class DateAround(Model):
    class Options:
        serialize_when_none = False

    origin = DateType(required=True)
    offset = StringType(regex=r"\d+d")
    scale = StringType(regex=r"\d+d")
    decay = FloatType()

class LocationAround(Model):
    class Options:
        serialize_when_none = False

    origin = StringType(regex=r"(-?\d+(\.\d+)?),(-?\d+(\.\d+)?)", required=True)
    offset = StringType(regex=r"\d+(\.\d+)?(cm|m|km|in|ft|mi)")
    scale = StringType(regex=r"\d+(\.\d+)?(cm|m|km|in|ft|mi)")
    decay = FloatType()

class DateTimeRange(Model):
    class Options:
        serialize_when_none = False

    gt = DateOrDateTimeType()
    gte = DateOrDateTimeType()
    lt = DateOrDateTimeType()
    lte = DateOrDateTimeType()
    tz = StringType(choices=pytz.all_timezones)

class DateRange(Model):
    class Options:
        serialize_when_none = False

    gt = DateType()
    gte = DateType()
    lt = DateType()
    lte = DateType()
    tz = StringType(choices=pytz.all_timezones)

class FloatRange(Model):
    class Options:
        serialize_when_none = False

    gt = FloatType()
    gte = FloatType()
    lt = FloatType()
    lte = FloatType()

class IntRange(Model):
    class Options:
        serialize_when_none = False

    gt = IntType()
    gte = IntType()
    lt = IntType()
    lte = IntType()

class BrandUnsafe(Model):
    class Options:
        serialize_when_none = False

    exclude = BooleanType()

class Entity(Model):
    class Options:
        serialize_when_none = False

    id = StringType()

class LimitMixin(Model):

    limit = IntType(min_value=1)

class OffsetMixin(Model):

    offset = IntType(min_value=0)

class PaginatedMixin(LimitMixin, OffsetMixin):

    pass

class SortableMixin(Model):

    sort = ListType(StringType())

class _Config(Model):

    verify_ssl = BooleanType(default=True)

class ConfigMixin(Model):

    config = ModelType(_Config)

class ResultSet(Model):

    count = IntType()
    previous = URLType()
    next = URLType()

    @property
    def results(self):  # pragma: nocover
        raise NotImplementedError()

    def _parse_params(self, url):
        return dict(parse_qsl(urlparse(url).query))

    def has_previous(self):
        return self.previous is not None

    def has_next(self):
        return self.next is not None

    def get_next(self):
        if not self.has_next() or not hasattr(self, "_more"):
            return
        params = self._parse_params(self.next)
        return self._more(**params)

    def get_previous(self):
        if not self.has_previous() or not hasattr(self, "_more"):
            return
        params = self._parse_params(self.previous)
        return self._more(**params)

    def iter_pages(self):
        next_page = self
        while next_page.has_next():
            next_page = next_page.get_next()
            yield next_page

    def iter_items(self):
        return iter(self.results)

    def iter_all(self):
        for item in self.iter_items():
            yield item
        for page in self.iter_pages():
            for item in page.iter_items():
                yield item

    def __iter__(self):
        return self.iter_items()

class ResultType(SchematicsListType):
    def __init__(self, field, *args, **kwargs):
        super(ResultType, self).__init__(ModelType(field), *args, **kwargs)
<commit_msg>fix area radius regex<commit_after>import re
from datetime import datetime, date, time as dt_time
from urllib.parse import parse_qsl, urlparse

import pytz
from dateutil.parser import parse as parse_date

from schematics.exceptions import (
    ValidationError as SchematicsValidationError,
    DataError as SchematicsDataError,
    ConversionError,
)
from schematics.models import Model
from schematics.transforms import Role
from schematics.types import (
    StringType,
    DateTimeType as SchematicsDateTimeType,
    IntType,
    FloatType,
    URLType,
    GeoPointType,
    BooleanType,
    DateType as SchematicsDateType,
)
from schematics.types.compound import ListType as SchematicsListType, ModelType, DictType, PolyModelType
from schematics.types.serializable import serializable

class DateTimeType(SchematicsDateTimeType):
    def to_native(self, value, context=None):
        if isinstance(value, datetime):
            return value
        return parse_date(value)

class DateType(SchematicsDateType):
    def to_native(self, value, context=None):
        if isinstance(value, date):
            return value
        return parse_date(value).date()

class DateOrDateTimeType(SchematicsDateTimeType):
    def validate_tz(self, value, context=None):
        if isinstance(value, datetime):
            return super().validate_tz(value, context)

    def to_primitive(self, value, context=None):
        if isinstance(value, datetime):
            return super().to_primitive(value, context)
        else:
            return value.strftime("%Y-%m-%d")

    def to_native(self, value, context=None):
        if isinstance(value, str):
            try:
                return datetime.strptime(value, "%Y-%m-%d").date()
            except ValueError:
                return parse_date(value)
        if isinstance(value, datetime):
            return value
        return value

class StringModelType(ModelType):
    def _convert(self, value, context):
        if isinstance(value, str):
            return self.model_class._convert(value, context=context)
        else:
            return super(StringModelType, self)._convert(value, context=context)

class StringModel(Model):
    @property
    def import_format(self):  # pragma: nocover
        raise NotImplementedError()

    @property
    def export_format(self):  # pragma: nocover
        raise NotImplementedError()

    def convert(self, raw_data, **kw):
        if isinstance(raw_data, str):
            try:
                raw_data = re.match(self.import_format, str(raw_data)).groupdict()
            except AttributeError:
                raise SchematicsValidationError(f"'{raw_data}' is not a valid format for {self.__class__.__name__}")
        return super(StringModel, self).convert(raw_data, **kw)

    def export(self, *args, **kwargs):
        return self.export_format.format(**super(StringModel, self).export(*args, **kwargs))

class StringListType(SchematicsListType):
    def __init__(self, field, separator=" ", **kwargs):
        self.separator = separator
        super(StringListType, self).__init__(field, **kwargs)

    def _coerce(self, value):
        if isinstance(value, str):
            return value.split(self.separator)
        elif isinstance(value, dict):
            return [value]
        else:
            return super(StringListType, self)._coerce(value)

    def _export(self, *args, **kwargs):
        data = super(StringListType, self)._export(*args, **kwargs)
        return self.separator.join(map(str, data))

class ListType(SchematicsListType):
    def _coerce(self, value):
        if not isinstance(value, (list, tuple)):
            return [value]
        else:
            return super(ListType, self)._coerce(value)

class GeoJSONPointType(GeoPointType):
    def _normalize(self, value):
        return tuple(reversed(super(GeoJSONPointType, self)._normalize(value)))

class Area(StringModel):

    import_format = r"(?P<radius>(\d+(\.\d+)?)(km|mi)|(\d+)(m|ft))@(?P<latitude>[\-\+]?\d+(\.\d+)?),(?P<longitude>[\-\+]?\d+(\.\d+)?)"
    export_format = "{radius}@{latitude},{longitude}"

    radius = StringType(regex=r"((\d+(\.\d+)?)(km|mi)|(\d+)(m|ft))", required=True)
    latitude = FloatType(required=True)
    longitude = FloatType(required=True)

class Location(StringModel):

    import_format = r"@(?P<latitude>-?\d+(\.\d+)?),(?P<longitude>-?\d+(\.\d+)?)"
    export_format = "@{latitude},{longitude}"

    latitude = FloatType(required=True)
    longitude = FloatType(required=True)

class Place(Model):
    class Options:
        serialize_when_none = False

    scope = ListType(StringType)
    exact = ListType(StringType)

class DateAround(Model):
    class Options:
        serialize_when_none = False

    origin = DateType(required=True)
    offset = StringType(regex=r"\d+d")
    scale = StringType(regex=r"\d+d")
    decay = FloatType()

class LocationAround(Model):
    class Options:
        serialize_when_none = False

    origin = StringType(regex=r"(-?\d+(\.\d+)?),(-?\d+(\.\d+)?)", required=True)
    offset = StringType(regex=r"\d+(\.\d+)?(cm|m|km|in|ft|mi)")
    scale = StringType(regex=r"\d+(\.\d+)?(cm|m|km|in|ft|mi)")
    decay = FloatType()

class DateTimeRange(Model):
    class Options:
        serialize_when_none = False

    gt = DateOrDateTimeType()
    gte = DateOrDateTimeType()
    lt = DateOrDateTimeType()
    lte = DateOrDateTimeType()
    tz = StringType(choices=pytz.all_timezones)

class DateRange(Model):
    class Options:
        serialize_when_none = False

    gt = DateType()
    gte = DateType()
    lt = DateType()
    lte = DateType()
    tz = StringType(choices=pytz.all_timezones)

class FloatRange(Model):
    class Options:
        serialize_when_none = False

    gt = FloatType()
    gte = FloatType()
    lt = FloatType()
    lte = FloatType()

class IntRange(Model):
    class Options:
        serialize_when_none = False

    gt = IntType()
    gte = IntType()
    lt = IntType()
    lte = IntType()

class BrandUnsafe(Model):
    class Options:
        serialize_when_none = False

    exclude = BooleanType()

class Entity(Model):
    class Options:
        serialize_when_none = False

    id = StringType()

class LimitMixin(Model):

    limit = IntType(min_value=1)

class OffsetMixin(Model):

    offset = IntType(min_value=0)

class PaginatedMixin(LimitMixin, OffsetMixin):

    pass

class SortableMixin(Model):

    sort = ListType(StringType())

class _Config(Model):

    verify_ssl = BooleanType(default=True)

class ConfigMixin(Model):

    config = ModelType(_Config)

class ResultSet(Model):

    count = IntType()
    previous = URLType()
    next = URLType()

    @property
    def results(self):  # pragma: nocover
        raise NotImplementedError()

    def _parse_params(self, url):
        return dict(parse_qsl(urlparse(url).query))

    def has_previous(self):
        return self.previous is not None

    def has_next(self):
        return self.next is not None

    def get_next(self):
        if not self.has_next() or not hasattr(self, "_more"):
            return
        params = self._parse_params(self.next)
        return self._more(**params)

    def get_previous(self):
        if not self.has_previous() or not hasattr(self, "_more"):
            return
        params = self._parse_params(self.previous)
        return self._more(**params)

    def iter_pages(self):
        next_page = self
        while next_page.has_next():
            next_page = next_page.get_next()
            yield next_page

    def iter_items(self):
        return iter(self.results)

    def iter_all(self):
        for item in self.iter_items():
            yield item
        for page in self.iter_pages():
            for item in page.iter_items():
                yield item

    def __iter__(self):
        return self.iter_items()

class ResultType(SchematicsListType):
    def __init__(self, field, *args, **kwargs):
        super(ResultType, self).__init__(ModelType(field), *args, **kwargs)
<|endoftext|>


<commit_before>import json
import time

from lxml import etree
import requests

EN_SONG_SEARCH_URL = 'http://developer.echonest.com/api/v4/song/search'

def found_song(xml_path, api_key):
  params = {
    'api_key': api_key,
    'format': 'json',
    'results': 1
  }
  try:
    root = etree.parse(xml_path)
  except:
    with open('error.txt', 'a') as f:
      f.write(os.path.basename(xml_path) + '\n')
    return

  artist = root.xpath('//artist/text()')
  if artist:
    params['artist'] = artist

  title = root.xpath('//title/text()')
  if title:
    params['title'] = title

  resp = requests.get(EN_SONG_SEARCH_URL, params=params)
  resp.raise_for_status()
  data = json.loads(resp.text)
  return bool(data['response']['status']['code'] == 0 and
              len(data['response']['songs']))

if __name__ == '__main__':
  import errno
  import os
  import sys
  import shutil
  import util

  util.mkdir_p('xml/existing')
  util.mkdir_p('xml/nonexistant')
  api_key = os.environ['ECHO_NEST_API_KEY']

  if len(sys.argv) != 2:
    # No argument given, do entire xml/raw directory
    for dirpath, dirnames, filenames in os.walk('xml/raw'):
      for filename in filenames:
        full_path = os.path.join(dirpath, filename)
        print full_path,

        exist_path = os.path.join('xml/existing', filename)
        if os.path.isfile(exist_path):
          print ' o'
          continue

        nonexist_path = os.path.join('xml/nonexistant', filename)
        if os.path.isfile(nonexist_path):
          print ' x'
          continue

        if found_song(full_path, api_key):
          print ' O'
          shutil.copyfile(full_path, exist_path)
        else:
          util.touch(nonexist_path)
          print ' X'
        time.sleep(0.5)
  else:
    if found_song(sys.argv[1], api_key):
      print os.path.join('xml/existing', os.path.basename(sys.argv[1]))
<commit_msg>MUCH stricter about what constitutes a "real" song.<commit_after>import json
import time

from lxml import etree
import requests

EN_SONG_SEARCH_URL = 'http://developer.echonest.com/api/v4/song/search'

def found_song(xml_path, api_key):
  params = {
    'api_key': api_key,
    'format': 'json',
    'results': 1
  }
  try:
    root = etree.parse(xml_path)
  except:
    with open('error.txt', 'a') as f:
      f.write(os.path.basename(xml_path) + '\n')
    return

  artist = root.xpath('//artist/text()')
  title = root.xpath('//title/text()')
  if not (artist and title):
    # Skip anything that doesn't have an artist AND a title
    return False

  params['artist'] = artist
  params['title'] = title

  resp = requests.get(EN_SONG_SEARCH_URL, params=params)
  resp.raise_for_status()
  data = json.loads(resp.text)
  return bool(data['response']['status']['code'] == 0 and
              len(data['response']['songs']))

if __name__ == '__main__':
  import errno
  import os
  import sys
  import shutil
  import util

  util.mkdir_p('xml/existing')
  util.mkdir_p('xml/nonexistant')
  api_key = os.environ['ECHO_NEST_API_KEY']

  if len(sys.argv) != 2:
    # No argument given, do entire xml/raw directory
    for dirpath, dirnames, filenames in os.walk('xml/raw'):
      for filename in filenames:
        full_path = os.path.join(dirpath, filename)
        print full_path,

        exist_path = os.path.join('xml/existing', filename)
        if os.path.isfile(exist_path):
          print ' o'
          continue

        nonexist_path = os.path.join('xml/nonexistant', filename)
        if os.path.isfile(nonexist_path):
          print ' x'
          continue

        if found_song(full_path, api_key):
          print ' O'
          shutil.copyfile(full_path, exist_path)
        else:
          util.touch(nonexist_path)
          print ' X'
        time.sleep(0.5)
  else:
    if found_song(sys.argv[1], api_key):
      print os.path.join('xml/existing', os.path.basename(sys.argv[1]))
<|endoftext|>


<commit_before># python stuff
import io
import os
import sys
import types
import traceback

# binja stuff
import binaryninjaui
from binaryninja.settings import Settings
from binaryninja import log
from binaryninja import _binaryninjacore as core
from binaryninjaui import View, ViewType, ViewFrame, UIContext, HexEditor
from binaryninja import binaryview
from PySide2.QtWidgets import QScrollArea, QWidget, QVBoxLayout, QGroupBox, QTreeWidget, QTreeWidgetItem, QLineEdit
from PySide2.QtCore import Qt

if sys.version_info[0] == 2:
	import kshelpers
else:
	from . import kshelpers

class KaitaiView(QScrollArea, View):
	def __init__(self, parent, binaryView):
		QScrollArea.__init__(self, parent)

		View.__init__(self)
		self.setupView(self)

		# BinaryViewType
		self.binaryView = binaryView

		self.rootSelectionStart = 0
		self.rootSelectionEnd = 1

		self.ioRoot = None
		self.ioCurrent = None

		container = QWidget(self)
		layout = QVBoxLayout()

		self.treeGroup = QGroupBox("Data Tree:")
		treeLayout = QVBoxLayout()
		self.treeWidget = QTreeWidget()
		self.treeWidget.setColumnCount(4)
		self.treeWidget.setHeaderLabels(['label','value','start','end'])
		self.treeWidget.itemSelectionChanged.connect(self.onTreeSelect)
		self.structPath = QLineEdit("root")
		self.structPath.setDisabled(True)
		treeLayout.addWidget(self.structPath)
		treeLayout.addWidget(self.treeWidget)
		self.treeGroup.setLayout(treeLayout)

		self.hexGroup = QGroupBox("Hex View:")
		self.hexLayout = QVBoxLayout()
		self.hexWidget = HexEditor(binaryView, ViewFrame.viewFrameForWidget(self), 0)
		self.hexLayout.addWidget(self.hexWidget)
		self.hexGroup.setLayout(self.hexLayout)

		layout.addWidget(self.treeGroup)
		layout.addWidget(self.hexGroup)
		#layout.addStretch(1)
		container.setLayout(layout)
		self.setWidgetResizable(True)
		self.setWidget(container)

		self.kaitaiParse()

	# parse the file using Kaitai, construct the TreeWidget
	def kaitaiParse(self):
		parsed = None

		kaitaiIO = kshelpers.KaitaiBinaryViewIO(self.binaryView)
		if not kaitaiIO:
			print('ERROR: initializing kaitai binary view')
		parsed = kshelpers.parseIo(kaitaiIO)
		if not parsed:
			print('ERROR: parsing the binary view')

		tree = kshelpers.buildQtree(parsed)
		if not tree:
			return

		self.ioRoot = tree.ksobj._io
		self.ioCurrent = tree.ksobj._io

		self.treeWidget.clear()
		self.treeWidget.setSortingEnabled(False)						# temporarily, for efficiency
		# two options with how we create the hierarchy
		if False:
			# treat root as top level "file" container
			tree.setLabel('file')
			tree.setValue(None)
			tree.setStart(0)
			tree.setEnd(0)
			self.treeWidget.insertTopLevelItem(0, tree)
		else:
			# add root's children as top level items
			self.treeWidget.insertTopLevelItems(0, tree.takeChildren())

		#self.treeWidget.expandAll()
		width = self.treeWidget.width()
		self.treeWidget.setColumnWidth(0, .4*width)
		self.treeWidget.setColumnWidth(1, .2*width)
		self.treeWidget.setColumnWidth(2, .1*width)
		self.treeWidget.setColumnWidth(3, .1*width)
		# enable sorting
		self.treeWidget.setSortingEnabled(True)
		self.treeWidget.sortByColumn(2, Qt.AscendingOrder)

		# TODO: select first item, maybe expand a few things
		self.rootSelectionStart = 0
		self.rootSelectionEnd = 1
		self.hexWidget.setSelectionRange(0,1)
		self.treeWidget.setUniformRowHeights(True)

	# callback!
	def getData(self):
		return self.binaryView

	# callback!
	def getCurrentOffset(self):
		middle = self.rootSelectionStart + int((self.rootSelectionEnd - self.rootSelectionStart)/2)
		return middle

	def setCurrentOffset(self, offset):
		print('setCurrentOffset(0x%X)' % offset)
		self.rootSelectionStart = offset
		UIContext.updateStatus(True)

	def getFont(self):
		return binaryninjaui.getMonospaceFont(self)

	def navigate(self, addr):
		return False

	def onTreeSelect(self, wtf=None):
		# get KaitaiTreeWidgetItem
		item = self.treeWidget.selectedItems()[0]

		# build path, inform user
		structPath = item.label
		itemTmp = item
		while itemTmp.parent():
			itemTmp = itemTmp.parent()
			structPath = itemTmp.label + '.' + structPath
		self.structPath.setText('root.' + structPath)

		#
		(start, end) = (item.start, item.end)
		if start == None or end == None:
			return

		# determine current IO we're in (the Kaitai input/output abstraction)
		_io = None
		# if the tree item is linked to a KaitaiNode, simply read the IO
		if item.ksobj:
			_io = item.ksobj._parent._io
		else:
			# else we're a leaf
			parent = item.parent()
			if parent:
				# a leaf with a parent -> read parent's IO
				_io = parent.ksobj._io
			else:
				# a leaf without a parent -> we must be at root -> use root IO
				_io = self.ioRoot

		# if the selection is in the root view, store the interval so that upon
		# getCurrentOffset() callback, we return the middle and feature map is
		# updated
		if _io == self.ioRoot:
			self.rootSelectionStart = start
			self.rootSelectionEnd = end

		# current kaitai object is on a different io? then swap HexEditor
		if _io != self.ioCurrent:
			# delete old view
			layoutItem = self.hexLayout.takeAt(0)
			hexEditorWidget = layoutItem.widget()
			hexEditorWidget.setParent(None)
			hexEditorWidget.deleteLater()
			self.hexWidget = None

			# if it's the original file IO, wrap the already-open file binary view
			if _io == self.ioRoot:
				self.hexWidget = HexEditor(self.binaryView, ViewFrame.viewFrameForWidget(self), 0)
			# otherwise delete old view, create a temporary view
			else:
				# create new view
				length = _io.size()
				_io.seek(0)
				data = _io.read_bytes(length)
				bv = binaryview.BinaryView.new(data)
				self.hexWidget = HexEditor(bv, ViewFrame.viewFrameForWidget(self), 0)

			self.hexLayout.addWidget(self.hexWidget)
			self.ioCurrent = _io

		# now position selection in whatever HexEditor is current
		#print('selecting to [0x%X, 0x%X)' % (start, end))
		self.hexWidget.setSelectionRange(start, end)

		# set hex group title to reflect current selection
		self.hexGroup.setTitle('Hex View @ [0x%X, 0x%X)' % (start, end))

class KaitaiViewType(ViewType):
	def __init__(self):
		super(KaitaiViewType, self).__init__("Kaitai", "Kaitai")

	# binaryView:		BinaryView
	def getPriority(self, binaryView, filename):
		# data.file:	FileMetadata
		# data.raw:		BinaryView

		priority = 0
		isExec = binaryView.executable
		weRecognize = kshelpers.idData(binaryView.read(0,16), len(binaryView)) != None

		# NOTE: ui/shared/hexeditor.cpp has the hex editor at priority 20 when
		# executable, and 10 otherwise

		if not weRecognize:
			priority = 0
		else:
			if isExec:
				priority = 21
			else:
				priority = 11

#		if isExec and weRecognize:
#			print('priority=25 to slightly beat out the hex editor case=(executable, recognize)')
#			priority = 25
#		if isExec and not weRecognize:
#			print('priority=15 to lose to the hex editor case=(executable, !recognize)')
#			priority = 15
#		if not isExec and weRecognize:
#			print('priority=100 to beat out everything case=(!executable, recognize)')
#			priority = 100
#		if not isExec and not weRecognize:
#			print('priority=5 to lose to hex editor case=(!executable, !recognize)')
#			priority = 5

		#print('returning priority=%d for KaitaiViewType' % priority)
		return priority

	def create(self, binaryView, view_frame):
		return KaitaiView(view_frame, binaryView)

ViewType.registerViewType(KaitaiViewType())
<commit_msg>kaitai: priority again, and dump the QGroupBoxes (take#2)<commit_after># python stuff
import io
import os
import sys
import types
import traceback

# binja stuff
import binaryninjaui
from binaryninja.settings import Settings
from binaryninja import log
from binaryninja import _binaryninjacore as core
from binaryninjaui import View, ViewType, ViewFrame, UIContext, HexEditor
from binaryninja import binaryview
from PySide2.QtWidgets import QScrollArea, QWidget, QVBoxLayout, QGroupBox, QTreeWidget, QTreeWidgetItem, QLineEdit, QHeaderView
from PySide2.QtCore import Qt

if sys.version_info[0] == 2:
	import kshelpers
else:
	from . import kshelpers

class MyQTreeWidget(QTreeWidget):
	def __init__(self):
		QTreeWidget.__init__(self)

	def resizeEvent(self, event):
		QTreeWidget.resizeEvent(self, event)
		self.scrollTo(self.currentIndex())

class KaitaiView(QScrollArea, View):
	def __init__(self, parent, binaryView):
		QScrollArea.__init__(self, parent)

		View.__init__(self)
		self.setupView(self)

		# BinaryViewType
		self.binaryView = binaryView

		self.rootSelectionStart = 0
		self.rootSelectionEnd = 1

		self.ioRoot = None
		self.ioCurrent = None

		container = QWidget(self)
		self.layout = QVBoxLayout()

		self.treeWidget = MyQTreeWidget()
		self.treeWidget.setColumnCount(4)
		self.treeWidget.setHeaderLabels(['label','value','start','end'])
		self.treeWidget.itemSelectionChanged.connect(self.onTreeSelect)

		self.structPath = QLineEdit("root")
		self.structPath.setDisabled(True)

		self.hexWidget = HexEditor(binaryView, ViewFrame.viewFrameForWidget(self), 0)

		self.layout.addWidget(self.treeWidget)
		self.layout.addWidget(self.structPath)
		self.layout.addWidget(self.hexWidget)
		#self.layout.addStretch(1)
		container.setLayout(self.layout)
		self.setWidgetResizable(True)
		self.setWidget(container)

		self.kaitaiParse()

	# parse the file using Kaitai, construct the TreeWidget
	def kaitaiParse(self):
		parsed = None

		kaitaiIO = kshelpers.KaitaiBinaryViewIO(self.binaryView)
		if not kaitaiIO:
			print('ERROR: initializing kaitai binary view')
		parsed = kshelpers.parseIo(kaitaiIO)
		if not parsed:
			print('ERROR: parsing the binary view')

		tree = kshelpers.buildQtree(parsed)
		if not tree:
			return

		self.ioRoot = tree.ksobj._io
		self.ioCurrent = tree.ksobj._io

		self.treeWidget.clear()
		self.treeWidget.setSortingEnabled(False)						# temporarily, for efficiency
		# two options with how we create the hierarchy
		if False:
			# treat root as top level "file" container
			tree.setLabel('file')
			tree.setValue(None)
			tree.setStart(0)
			tree.setEnd(0)
			self.treeWidget.insertTopLevelItem(0, tree)
		else:
			# add root's children as top level items
			self.treeWidget.insertTopLevelItems(0, tree.takeChildren())

		# enable sorting
		self.treeWidget.setSortingEnabled(True)
		self.treeWidget.sortByColumn(2, Qt.AscendingOrder)

		# TODO: select first item, maybe expand a few things
		self.rootSelectionStart = 0
		self.rootSelectionEnd = 1
		self.hexWidget.setSelectionRange(0,1)
		self.treeWidget.setUniformRowHeights(True)

	# Qt callbacks
	def resizeEvent(self, event):
		width = self.treeWidget.width()
		self.treeWidget.setColumnWidth(0, .6*width)
		self.treeWidget.setColumnWidth(1, .2*width)
		self.treeWidget.setColumnWidth(2, .1*width)
		self.treeWidget.setColumnWidth(3, .1*width)

		QScrollArea.resizeEvent(self, event)

	# binja callbacks
	def getData(self):
		return self.binaryView

	def getStart(self):
		result = self.binaryView.start
		#print('getStart() returning '+str(result))
		return result

	def getEnd(self):
		result = self.binaryView.end
		#print('getEnd() returning '+str(result))
		return result

	def getLength(self):
		result = len(self.binaryView)
		#print('getLength() returning '+str(result))
		return result

	def getCurrentOffset(self):
		result = self.rootSelectionStart + int((self.rootSelectionEnd - self.rootSelectionStart)/2)
		#result = self.rootSelectionStart
		#print('getCurrentOffset() returning '+str(result))
		return result

	def getSelectionOffsets(self):
		result = None
		if self.hexWidget:
			result = self.hexWidget.getSelectionOffsets()
		else:
			result = (self.rootSelectionStart, self.rootSelectionStart)
		#print('getSelectionOffsets() returning '+str(result))
		return result

	def setCurrentOffset(self, offset):
		#print('setCurrentOffset(0x%X)' % offset)
		self.rootSelectionStart = offset
		UIContext.updateStatus(True)

	def getFont(self):
		return binaryninjaui.getMonospaceFont(self)

	def navigate(self, addr):
		#print('navigate()')
		return False

	def navigateToFileOffset(self, offset):
		#print('navigateToFileOffset()')
		return False

	def onTreeSelect(self, wtf=None):
		# get KaitaiTreeWidgetItem
		item = self.treeWidget.selectedItems()[0]

		# build path, inform user
		structPath = item.label
		itemTmp = item
		while itemTmp.parent():
			itemTmp = itemTmp.parent()
			structPath = itemTmp.label + '.' + structPath
		self.structPath.setText('root.' + structPath)

		#
		(start, end) = (item.start, item.end)
		if start == None or end == None:
			return

		# determine current IO we're in (the Kaitai input/output abstraction)
		_io = None
		# if the tree item is linked to a KaitaiNode, simply read the IO
		if item.ksobj:
			_io = item.ksobj._parent._io
		else:
			# else we're a leaf
			parent = item.parent()
			if parent:
				# a leaf with a parent -> read parent's IO
				_io = parent.ksobj._io
			else:
				# a leaf without a parent -> we must be at root -> use root IO
				_io = self.ioRoot

		# if the selection is in the root view, store the interval so that upon
		# getCurrentOffset() callback, we return the middle and feature map is
		# updated
		if _io == self.ioRoot:
			self.rootSelectionStart = start
			self.rootSelectionEnd = end

		# current kaitai object is on a different io? then swap HexEditor
		if _io != self.ioCurrent:
			# delete old view
			layoutItem = self.layout.takeAt(2)
			hexEditorWidget = layoutItem.widget()
			hexEditorWidget.setParent(None)
			hexEditorWidget.deleteLater()
			self.hexWidget = None

			# if it's the original file IO, wrap the already-open file binary view
			if _io == self.ioRoot:
				self.hexWidget = HexEditor(self.binaryView, ViewFrame.viewFrameForWidget(self), 0)
			# otherwise delete old view, create a temporary view
			else:
				# create new view
				length = _io.size()
				_io.seek(0)
				data = _io.read_bytes(length)
				bv = binaryview.BinaryView.new(data)
				self.hexWidget = HexEditor(bv, ViewFrame.viewFrameForWidget(self), 0)

			self.layout.addWidget(self.hexWidget)
			self.ioCurrent = _io

		# now position selection in whatever HexEditor is current
		#print('selecting to [0x%X, 0x%X)' % (start, end))
		self.hexWidget.setSelectionRange(start, end)

		# set hex group title to reflect current selection
		#self.hexGroup.setTitle('Hex View @ [0x%X, 0x%X)' % (start, end))

class KaitaiViewType(ViewType):
	def __init__(self):
		super(KaitaiViewType, self).__init__("Kaitai", "Kaitai")

	# binaryView:		BinaryView
	def getPriority(self, binaryView, filename):
		#return 100
		#print('len(bv)=0x%X executable=%d bytes=%s' % (len(binaryView), binaryView.executable, repr(binaryView.read(0,4))))

		# executable means the view is mapped like an OS loader would load an executable (eg: view=ELF)
		# !executable means executable image is not mapped (eg: view=Raw) (or something like .png is loaded)
		if binaryView.executable:
			return 0

		if binaryView.start != 0:
			return 0

		if os.path.getsize(binaryView.file.filename) != len(binaryView):
			return 0

		dataSample = binaryView.read(0, 16)
		if len(dataSample) != 16:
			return 0

		# if we don't recognize it, return 0
		ksModuleName = kshelpers.idData(dataSample, len(binaryView))
		if not ksModuleName:
			return 0

		# for executables, yield triage (25)
		if ksModuleName in ['elf', 'microsoft_pe', 'mach_o']:
			return 24

		#
		return 100

	def create(self, binaryView, view_frame):
		return KaitaiView(view_frame, binaryView)

ViewType.registerViewType(KaitaiViewType())
<|endoftext|>


<commit_before>import plpy
import presto_client
from collections import namedtuple
import time

# Maximum length for identifiers (e.g. table names, column names, function names)
# defined in pg_config_manual.h
PG_NAMEDATALEN = 64

# convert Presto query result field types to PostgreSQL types
def _pg_result_type(presto_type):
    if presto_type == "varchar":
        return "varchar(255)"
    if presto_type == "varbinary":
        return "bytea"
    elif presto_type == "double":
        return "double precision"
    elif presto_type == "row":
        return "json"  # TODO record or anyarray???
    else:
        # assuming Presto and PostgreSQL use the same SQL standard name
        return presto_type

# convert Presto table column types to PostgreSQL types
def _pg_table_type(presto_type):
    if presto_type == "varchar":
        return "varchar(255)"
    if presto_type == "varbinary":
        return "bytea"
    elif presto_type == "double":
        return "double precision"
    elif presto_type == "row":
        return "json"  # TODO record or anyarray???
    else:
        # assuming Presto and PostgreSQL use the same SQL standard name
        return presto_type

# queries can include same column name twice but tables can't.
def _rename_duplicated_column_names(column_names):
    renamed = []
    used_names = set()
    for original_name in column_names:
        name = original_name
        while name in used_names:
            name += "_"
        if name != original_name:
            plpy.warning("Result column %s is renamed to %s because the name appears twice in a query result" % \
                    (plpy.quote_ident(original_name), plpy.quote_ident(name)))
        used_names.add(name)
        renamed.append(name)
    return renamed

# build CREATE TEMPORARY TABLE statement
def _build_create_temp_table_sql(table_name, column_names, column_types):
    create_sql = ["create temporary table %s (\n  " % plpy.quote_ident(table_name)]

    first = True
    for column_name, column_type in zip(column_names, column_types):
        if first:
            first = False
        else:
            create_sql.append(",\n  ")

        create_sql.append(plpy.quote_ident(column_name))
        create_sql.append(" ")
        create_sql.append(column_type)

    create_sql.append("\n)")
    return ''.join(create_sql)

# build CREATE TABLE statement
def _build_create_table(schema_name, table_name, column_names, column_types, not_nulls):
    alter_sql = ["create table %s.%s (\n  " % (plpy.quote_ident(schema_name), plpy.quote_ident(table_name))]

    first = True
    for column_name, column_type, not_null in zip(column_names, column_types, not_nulls):
        if first:
            first = False
        else:
            alter_sql.append(",\n  ")

        alter_sql.append("%s %s" % (plpy.quote_ident(column_name), column_type))

        if not_null:
            alter_sql.append(" not null")

    alter_sql.append("\n)")
    return ''.join(alter_sql)

def _get_session_time_zone():
    rows = plpy.execute("show timezone")
    return rows[0].values()[0]

def _get_session_search_path_array():
    rows = plpy.execute("select ('{' || current_setting('search_path') || '}')::text[]")
    return rows[0].values()[0]

class QueryAutoClose(object):
    def __init__(self, query):
        self.query = query

    def __del__(self):
        self.query.close()

class QueryAutoCloseIterator(object):
    def __init__(self, gen, query_auto_close):
        self.gen = gen
        self.query_auto_close = query_auto_close

    def __iter__(self):
        return self

    def next(self):
        row = next(self.gen)
        return row

class SessionData(object):
    def __init__(self):
        self.query_auto_close = None

session = SessionData()

def start_presto_query(presto_server, presto_user, presto_catalog, presto_schema, function_name, query):
    try:
        # preserve search_path if explicitly set
        search_path = _get_session_search_path_array()
        if search_path != ['$user', 'public'] and len(search_path) > 0:
            # search_path is changed explicitly. use the first schema
            presto_schema = search_path[0]

        # start query
        client = presto_client.Client(server=presto_server, user=presto_user, catalog=presto_catalog, schema=presto_schema, time_zone=_get_session_time_zone())

        query = client.query(query)
        session.query_auto_close = QueryAutoClose(query)
        try:
            # result schema
            column_names = []
            column_types = []
            for column in query.columns():
                column_names.append(column.name)
                column_types.append(_pg_result_type(column.type))

            column_names = _rename_duplicated_column_names(column_names)

            # CREATE TABLE for return type of the function
            type_name = function_name + "_type"
            create_type_sql = _build_create_temp_table_sql(type_name, column_names, column_types)

            # CREATE FUNCTION
            create_function_sql = \
                """
                create or replace function pg_temp.%s()
                returns setof pg_temp.%s as $$
                    import prestogres
                    return prestogres.fetch_presto_query_results()
                $$ language plpythonu
                """ % \
                (plpy.quote_ident(function_name), plpy.quote_ident(type_name))

            # run statements
            plpy.execute("drop table if exists pg_temp.%s cascade" % \
                    (plpy.quote_ident(type_name)))
            plpy.execute(create_type_sql)
            plpy.execute(create_function_sql)

            query = None

        finally:
            if query is not None:
                # close query
                session.query_auto_close = None

    except (plpy.SPIError, presto_client.PrestoException) as e:
        # PL/Python converts an exception object in Python to an error message in PostgreSQL
        # using exception class name if exc.__module__ is either of "builtins", "exceptions",
        # or "__main__". Otherwise using "module.name" format. Set __module__ = "__module__"
        # to generate pretty messages.
        e.__class__.__module__ = "__main__"
        raise

def fetch_presto_query_results():
    try:
        # TODO should throw an exception?
        #if session.query_auto_close is None:

        results = session.query_auto_close.query.results()

        # close of the iterator closes query
        ite = QueryAutoCloseIterator(results, session.query_auto_close)
        session.query_auto_close = None
        return ite

    except (plpy.SPIError, presto_client.PrestoException) as e:
        e.__class__.__module__ = "__main__"
        raise

Column = namedtuple("Column", ("name", "type", "nullable"))

def setup_system_catalog(presto_server, presto_user, presto_catalog, access_role):
    client = presto_client.Client(server=presto_server, user=presto_user, catalog=presto_catalog, schema='default')

    # get table list
    sql = "select table_schema, table_name, column_name, is_nullable, data_type" \
          " from information_schema.columns"
    columns, rows = client.run(sql)
    if rows is None:
        rows = []

    schemas = {}

    for row in rows:
        schema_name = row[0]
        table_name = row[1]
        column_name = row[2]
        is_nullable = row[3]
        column_type = row[4]

        if schema_name == "sys" or schema_name == "information_schema":
            # skip system schemas
            continue

        if len(schema_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Schema %s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), PG_NAMEDATALEN - 1))
            continue

        tables = schemas.setdefault(schema_name, {})

        if len(table_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Table %s.%s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), plpy.quote_ident(table_name), PG_NAMEDATALEN - 1))
            continue

        columns = tables.setdefault(table_name, [])

        if len(column_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Column %s.%s.%s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), plpy.quote_ident(table_name), \
                     plpy.quote_ident(column_name), PG_NAMEDATALEN - 1))
            continue

        columns.append(Column(column_name, column_type, is_nullable))

    # drop all schemas excepting prestogres_catalog, information_schema and pg_%
    sql = "select n.nspname as schema_name from pg_catalog.pg_namespace n" \
          " where n.nspname not in ('prestogres_catalog', 'information_schema')" \
          " and n.nspname not like 'pg_%'"
    for row in plpy.cursor(sql):
        plpy.execute("drop schema %s cascade" % plpy.quote_ident(row["schema_name"]))

    # create schema and tables
    for schema_name, tables in sorted(schemas.items(), key=lambda (k,v): k):
        try:
            plpy.execute("create schema %s" % (plpy.quote_ident(schema_name)))
        except:
            # ignore error?
            pass

        # grant access on the all tables to the restricted user
        plpy.execute("grant select on all tables in schema %s to %s" % \
                (plpy.quote_ident(schema_name), plpy.quote_ident(access_role)))

        for table_name, columns in sorted(tables.items(), key=lambda (k,v): k):
            column_names = []
            column_types = []
            not_nulls = []
            for column in columns:
                column_names.append(column.name)
                column_types.append(_pg_table_type(column.type))
                not_nulls.append(not column.nullable)

            # change columns
            create_sql = _build_create_table(schema_name, table_name, column_names, column_types, not_nulls)
            plpy.execute(create_sql)

    # update pg_database
    plpy.execute("update pg_database set datname=%s where datname=current_database()" % \
            plpy.quote_literal(presto_catalog))

<commit_msg>fixed fetch_presto_query_results so that canceling query works<commit_after>import plpy
import presto_client
from collections import namedtuple
import time

# Maximum length for identifiers (e.g. table names, column names, function names)
# defined in pg_config_manual.h
PG_NAMEDATALEN = 64

# convert Presto query result field types to PostgreSQL types
def _pg_result_type(presto_type):
    if presto_type == "varchar":
        return "varchar(255)"
    if presto_type == "varbinary":
        return "bytea"
    elif presto_type == "double":
        return "double precision"
    elif presto_type == "row":
        return "json"  # TODO record or anyarray???
    else:
        # assuming Presto and PostgreSQL use the same SQL standard name
        return presto_type

# convert Presto table column types to PostgreSQL types
def _pg_table_type(presto_type):
    if presto_type == "varchar":
        return "varchar(255)"
    if presto_type == "varbinary":
        return "bytea"
    elif presto_type == "double":
        return "double precision"
    elif presto_type == "row":
        return "json"  # TODO record or anyarray???
    else:
        # assuming Presto and PostgreSQL use the same SQL standard name
        return presto_type

# queries can include same column name twice but tables can't.
def _rename_duplicated_column_names(column_names):
    renamed = []
    used_names = set()
    for original_name in column_names:
        name = original_name
        while name in used_names:
            name += "_"
        if name != original_name:
            plpy.warning("Result column %s is renamed to %s because the name appears twice in a query result" % \
                    (plpy.quote_ident(original_name), plpy.quote_ident(name)))
        used_names.add(name)
        renamed.append(name)
    return renamed

# build CREATE TEMPORARY TABLE statement
def _build_create_temp_table_sql(table_name, column_names, column_types):
    create_sql = ["create temporary table %s (\n  " % plpy.quote_ident(table_name)]

    first = True
    for column_name, column_type in zip(column_names, column_types):
        if first:
            first = False
        else:
            create_sql.append(",\n  ")

        create_sql.append(plpy.quote_ident(column_name))
        create_sql.append(" ")
        create_sql.append(column_type)

    create_sql.append("\n)")
    return ''.join(create_sql)

# build CREATE TABLE statement
def _build_create_table(schema_name, table_name, column_names, column_types, not_nulls):
    alter_sql = ["create table %s.%s (\n  " % (plpy.quote_ident(schema_name), plpy.quote_ident(table_name))]

    first = True
    for column_name, column_type, not_null in zip(column_names, column_types, not_nulls):
        if first:
            first = False
        else:
            alter_sql.append(",\n  ")

        alter_sql.append("%s %s" % (plpy.quote_ident(column_name), column_type))

        if not_null:
            alter_sql.append(" not null")

    alter_sql.append("\n)")
    return ''.join(alter_sql)

def _get_session_time_zone():
    rows = plpy.execute("show timezone")
    return rows[0].values()[0]

def _get_session_search_path_array():
    rows = plpy.execute("select ('{' || current_setting('search_path') || '}')::text[]")
    return rows[0].values()[0]

class QueryAutoClose(object):
    def __init__(self, query):
        self.query = query

    def __del__(self):
        self.query.close()

class QueryAutoCloseIterator(object):
    def __init__(self, gen, query_auto_close):
        self.gen = gen
        self.query_auto_close = query_auto_close

    def __iter__(self):
        return self

    def next(self):
        row = next(self.gen)
        return row

class SessionData(object):
    def __init__(self):
        self.query_auto_close = None

session = SessionData()

def start_presto_query(presto_server, presto_user, presto_catalog, presto_schema, function_name, query):
    try:
        # preserve search_path if explicitly set
        search_path = _get_session_search_path_array()
        if search_path != ['$user', 'public'] and len(search_path) > 0:
            # search_path is changed explicitly. use the first schema
            presto_schema = search_path[0]

        # start query
        client = presto_client.Client(server=presto_server, user=presto_user, catalog=presto_catalog, schema=presto_schema, time_zone=_get_session_time_zone())

        query = client.query(query)
        session.query_auto_close = QueryAutoClose(query)
        try:
            # result schema
            column_names = []
            column_types = []
            for column in query.columns():
                column_names.append(column.name)
                column_types.append(_pg_result_type(column.type))

            column_names = _rename_duplicated_column_names(column_names)

            # CREATE TABLE for return type of the function
            type_name = function_name + "_type"
            create_type_sql = _build_create_temp_table_sql(type_name, column_names, column_types)

            # CREATE FUNCTION
            create_function_sql = \
                """
                create or replace function pg_temp.%s()
                returns setof pg_temp.%s as $$
                    import prestogres
                    return prestogres.fetch_presto_query_results()
                $$ language plpythonu
                """ % \
                (plpy.quote_ident(function_name), plpy.quote_ident(type_name))

            # run statements
            plpy.execute("drop table if exists pg_temp.%s cascade" % \
                    (plpy.quote_ident(type_name)))
            plpy.execute(create_type_sql)
            plpy.execute(create_function_sql)

            query = None

        finally:
            if query is not None:
                # close query
                session.query_auto_close = None

    except (plpy.SPIError, presto_client.PrestoException) as e:
        # PL/Python converts an exception object in Python to an error message in PostgreSQL
        # using exception class name if exc.__module__ is either of "builtins", "exceptions",
        # or "__main__". Otherwise using "module.name" format. Set __module__ = "__module__"
        # to generate pretty messages.
        e.__class__.__module__ = "__main__"
        raise

def fetch_presto_query_results():
    try:
        # TODO should throw an exception?
        #if session.query_auto_close is None:

        query_auto_close = session.query_auto_close
        session.query_auto_close = None

        results = query_auto_close.query.results()

        # close of the iterator closes query
        return QueryAutoCloseIterator(results, query_auto_close)

    except (plpy.SPIError, presto_client.PrestoException) as e:
        e.__class__.__module__ = "__main__"
        raise

Column = namedtuple("Column", ("name", "type", "nullable"))

def setup_system_catalog(presto_server, presto_user, presto_catalog, access_role):
    client = presto_client.Client(server=presto_server, user=presto_user, catalog=presto_catalog, schema='default')

    # get table list
    sql = "select table_schema, table_name, column_name, is_nullable, data_type" \
          " from information_schema.columns"
    columns, rows = client.run(sql)
    if rows is None:
        rows = []

    schemas = {}

    for row in rows:
        schema_name = row[0]
        table_name = row[1]
        column_name = row[2]
        is_nullable = row[3]
        column_type = row[4]

        if schema_name == "sys" or schema_name == "information_schema":
            # skip system schemas
            continue

        if len(schema_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Schema %s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), PG_NAMEDATALEN - 1))
            continue

        tables = schemas.setdefault(schema_name, {})

        if len(table_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Table %s.%s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), plpy.quote_ident(table_name), PG_NAMEDATALEN - 1))
            continue

        columns = tables.setdefault(table_name, [])

        if len(column_name) > PG_NAMEDATALEN - 1:
            plpy.warning("Column %s.%s.%s is skipped because its name is longer than %d characters" % \
                    (plpy.quote_ident(schema_name), plpy.quote_ident(table_name), \
                     plpy.quote_ident(column_name), PG_NAMEDATALEN - 1))
            continue

        columns.append(Column(column_name, column_type, is_nullable))

    # drop all schemas excepting prestogres_catalog, information_schema and pg_%
    sql = "select n.nspname as schema_name from pg_catalog.pg_namespace n" \
          " where n.nspname not in ('prestogres_catalog', 'information_schema')" \
          " and n.nspname not like 'pg_%'"
    for row in plpy.cursor(sql):
        plpy.execute("drop schema %s cascade" % plpy.quote_ident(row["schema_name"]))

    # create schema and tables
    for schema_name, tables in sorted(schemas.items(), key=lambda (k,v): k):
        try:
            plpy.execute("create schema %s" % (plpy.quote_ident(schema_name)))
        except:
            # ignore error?
            pass

        # grant access on the all tables to the restricted user
        plpy.execute("grant select on all tables in schema %s to %s" % \
                (plpy.quote_ident(schema_name), plpy.quote_ident(access_role)))

        for table_name, columns in sorted(tables.items(), key=lambda (k,v): k):
            column_names = []
            column_types = []
            not_nulls = []
            for column in columns:
                column_names.append(column.name)
                column_types.append(_pg_table_type(column.type))
                not_nulls.append(not column.nullable)

            # change columns
            create_sql = _build_create_table(schema_name, table_name, column_names, column_types, not_nulls)
            plpy.execute(create_sql)

    # update pg_database
    plpy.execute("update pg_database set datname=%s where datname=current_database()" % \
            plpy.quote_literal(presto_catalog))

<|endoftext|>


<commit_before><commit_msg>heuristics<commit_after><|endoftext|>