<commit_before>import logging
from gi.repository import Gst

from lib.config import Config
from lib.tcpmulticonnection import TCPMultiConnection
from lib.clock import Clock

class AVPreviewOutput(TCPMultiConnection):

    def __init__(self, channel, port):
        self.log = logging.getLogger('AVPreviewOutput[{}]'.format(channel))
        super().__init__(port)

        self.channel = channel

        if Config.has_option('previews', 'videocaps'):
            target_caps = Config.get('previews', 'videocaps')
        else:
            target_caps = Config.get('mix', 'videocaps')

        pipeline = """
            intervideosrc channel=video_{channel} !
            {vcaps} !
            {vpipeline} !
            queue !
            mux.

            interaudiosrc channel=audio_{channel} !
            {acaps} !
            queue !
            mux.

            matroskamux
                name=mux
                streamable=true
                writing-app=Voctomix-AVPreviewOutput !

            multifdsink
                blocksize=1048576
                buffers-max=500
                sync-method=next-keyframe
                name=fd
        """.format(
            channel=self.channel,
            acaps=Config.get('mix', 'audiocaps'),
            vcaps=Config.get('mix', 'videocaps'),
            vpipeline=self.construct_video_pipeline(target_caps)
        )

        self.log.debug('Creating Output-Pipeline:\n%s', pipeline)
        self.outputPipeline = Gst.parse_launch(pipeline)
        self.outputPipeline.use_clock(Clock)

        self.log.debug('Binding Error & End-of-Stream-Signal '
                       'on Output-Pipeline')
        self.outputPipeline.bus.add_signal_watch()
        self.outputPipeline.bus.connect("message::eos", self.on_eos)
        self.outputPipeline.bus.connect("message::error", self.on_error)

        self.log.debug('Launching Output-Pipeline')
        self.outputPipeline.set_state(Gst.State.PLAYING)

    def construct_video_pipeline(self, target_caps):
        vaapi_enabled = Config.has_option('previews', 'vaapi')
        if vaapi_enabled:
            return self.construct_vaapi_video_pipeline(target_caps)

        else:
            return self.construct_native_video_pipeline(target_caps)

    def construct_vaapi_video_pipeline(self, target_caps):
        if Gst.version() < (1, 8):
            vaapi_encoders = {
                'h264': 'vaapiencode_h264',
                'jpeg': 'vaapiencode_jpeg',
                'mpeg2': 'vaapiencode_mpeg2',
            }
        else:
            vaapi_encoders = {
                'h264': 'vaapih264enc',
                'jpeg': 'vaapijpegenc',
                'mpeg2': 'vaapimpeg2enc',
            }

        vaapi_encoder_options = {
            'h264': 'rate-control=cqp init-qp=10 '
                    'max-bframes=0 keyframe-period=60',
            'jpeg': 'vaapiencode_jpeg quality=90'
                    'keyframe-period=0',
            'mpeg2': 'keyframe-period=60',
        }

        encoder = Config.get('previews', 'vaapi')
        do_deinterlace = Config.getboolean('previews', 'deinterlace')

        caps = Gst.Caps.from_string(target_caps)
        struct = caps.get_structure(0)
        _, width = struct.get_int('width')
        _, height = struct.get_int('height')
        _, framerate_numerator, framerate_denominator = struct.get_fraction('framerate')

        return '''
            capsfilter caps=video/x-raw,interlace-mode=progressive !
            vaapipostproc
                format=i420
                deinterlace-mode={imode}
                deinterlace-method=motion-adaptive
                width={width}
                height={height} !
            capssetter caps=video/x-raw,framerate={n}/{d} !
            {encoder} {options}
        '''.format(
            imode='interlaced' if do_deinterlace else 'disabled',
            width=width,
            height=height,
            encoder=vaapi_encoders[encoder],
            options=vaapi_encoder_options[encoder],
            n=framerate_numerator,
            d=framerate_denominator,
        )

    def construct_native_video_pipeline(self, target_caps):
        do_deinterlace = Config.getboolean('previews', 'deinterlace')

        return '''
            videoscale !
            {target_caps} !
            deinterlace mode={imode} !
            jpegenc quality=90
        '''.format(
            imode='interlaced' if do_deinterlace else 'disabled',
            target_caps=target_caps,
        )

    def on_accepted(self, conn, addr):
        self.log.debug('Adding fd %u to multifdsink', conn.fileno())
        fdsink = self.outputPipeline.get_by_name('fd')
        fdsink.emit('add', conn.fileno())

        def on_disconnect(multifdsink, fileno):
            if fileno == conn.fileno():
                self.log.debug('fd %u removed from multifdsink', fileno)
                self.close_connection(conn)

        fdsink.connect('client-fd-removed', on_disconnect)

    def on_eos(self, bus, message):
        self.log.debug('Received End-of-Stream-Signal on Output-Pipeline')

    def on_error(self, bus, message):
        self.log.debug('Received Error-Signal on Output-Pipeline')
        (error, debug) = message.parse_error()
        self.log.debug('Error-Details: #%u: %s', error.code, debug)
<commit_msg>fix non-vaapi-pipeline<commit_after>import logging
from gi.repository import Gst

from lib.config import Config
from lib.tcpmulticonnection import TCPMultiConnection
from lib.clock import Clock

class AVPreviewOutput(TCPMultiConnection):

    def __init__(self, channel, port):
        self.log = logging.getLogger('AVPreviewOutput[{}]'.format(channel))
        super().__init__(port)

        self.channel = channel

        if Config.has_option('previews', 'videocaps'):
            target_caps = Config.get('previews', 'videocaps')
        else:
            target_caps = Config.get('mix', 'videocaps')

        pipeline = """
            intervideosrc channel=video_{channel} !
            {vcaps} !
            {vpipeline} !
            queue !
            mux.

            interaudiosrc channel=audio_{channel} !
            {acaps} !
            queue !
            mux.

            matroskamux
                name=mux
                streamable=true
                writing-app=Voctomix-AVPreviewOutput !

            multifdsink
                blocksize=1048576
                buffers-max=500
                sync-method=next-keyframe
                name=fd
        """.format(
            channel=self.channel,
            acaps=Config.get('mix', 'audiocaps'),
            vcaps=Config.get('mix', 'videocaps'),
            vpipeline=self.construct_video_pipeline(target_caps)
        )

        self.log.debug('Creating Output-Pipeline:\n%s', pipeline)
        self.outputPipeline = Gst.parse_launch(pipeline)
        self.outputPipeline.use_clock(Clock)

        self.log.debug('Binding Error & End-of-Stream-Signal '
                       'on Output-Pipeline')
        self.outputPipeline.bus.add_signal_watch()
        self.outputPipeline.bus.connect("message::eos", self.on_eos)
        self.outputPipeline.bus.connect("message::error", self.on_error)

        self.log.debug('Launching Output-Pipeline')
        self.outputPipeline.set_state(Gst.State.PLAYING)

    def construct_video_pipeline(self, target_caps):
        vaapi_enabled = Config.has_option('previews', 'vaapi')
        if vaapi_enabled:
            return self.construct_vaapi_video_pipeline(target_caps)

        else:
            return self.construct_native_video_pipeline(target_caps)

    def construct_vaapi_video_pipeline(self, target_caps):
        if Gst.version() < (1, 8):
            vaapi_encoders = {
                'h264': 'vaapiencode_h264',
                'jpeg': 'vaapiencode_jpeg',
                'mpeg2': 'vaapiencode_mpeg2',
            }
        else:
            vaapi_encoders = {
                'h264': 'vaapih264enc',
                'jpeg': 'vaapijpegenc',
                'mpeg2': 'vaapimpeg2enc',
            }

        vaapi_encoder_options = {
            'h264': 'rate-control=cqp init-qp=10 '
                    'max-bframes=0 keyframe-period=60',
            'jpeg': 'vaapiencode_jpeg quality=90'
                    'keyframe-period=0',
            'mpeg2': 'keyframe-period=60',
        }

        encoder = Config.get('previews', 'vaapi')
        do_deinterlace = Config.getboolean('previews', 'deinterlace')

        caps = Gst.Caps.from_string(target_caps)
        struct = caps.get_structure(0)
        _, width = struct.get_int('width')
        _, height = struct.get_int('height')
        _, framerate_numerator, framerate_denominator = struct.get_fraction('framerate')

        return '''
            capsfilter caps=video/x-raw,interlace-mode=progressive !
            vaapipostproc
                format=i420
                deinterlace-mode={imode}
                deinterlace-method=motion-adaptive
                width={width}
                height={height} !
            capssetter caps=video/x-raw,framerate={n}/{d} !
            {encoder} {options}
        '''.format(
            imode='interlaced' if do_deinterlace else 'disabled',
            width=width,
            height=height,
            encoder=vaapi_encoders[encoder],
            options=vaapi_encoder_options[encoder],
            n=framerate_numerator,
            d=framerate_denominator,
        )

    def construct_native_video_pipeline(self, target_caps):
        do_deinterlace = Config.getboolean('previews', 'deinterlace')

        if do_deinterlace:
            pipeline = '''
                deinterlace mode={imode} !
                videorate !
            '''
        else:
            pipeline = ''

        pipeline += '''
            videoscale !
            {target_caps} !
            jpegenc quality=90
        '''

        return pipeline.format(
            imode='interlaced' if do_deinterlace else 'disabled',
            target_caps=target_caps,
        )

    def on_accepted(self, conn, addr):
        self.log.debug('Adding fd %u to multifdsink', conn.fileno())
        fdsink = self.outputPipeline.get_by_name('fd')
        fdsink.emit('add', conn.fileno())

        def on_disconnect(multifdsink, fileno):
            if fileno == conn.fileno():
                self.log.debug('fd %u removed from multifdsink', fileno)
                self.close_connection(conn)

        fdsink.connect('client-fd-removed', on_disconnect)

    def on_eos(self, bus, message):
        self.log.debug('Received End-of-Stream-Signal on Output-Pipeline')

    def on_error(self, bus, message):
        self.log.debug('Received Error-Signal on Output-Pipeline')
        (error, debug) = message.parse_error()
        self.log.debug('Error-Details: #%u: %s', error.code, debug)
<|endoftext|>


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

import os

from .markuputils import html_escape, xml_escape, attribute_escape

class _MarkupWriter(object):

    def __init__(self, output, line_separator=os.linesep, encoding=None):
        if isinstance(output, basestring):
            output = open(output, 'w')
        self.output = output
        self._line_separator = line_separator
        self._encode = self._create_encoder(encoding)

    def _create_encoder(self, encoding):
        if encoding is None:
            return lambda text: text
        return lambda text: text.encode(encoding)

    def start(self, name, attrs=None, newline=True):
        self._write('<%s%s>%s' % (name, self._get_attrs(attrs),
                                  self._line_separator if newline else ''))

    def content(self, content=None, escape=True):
        if content:
            self._write(self._escape(content) if escape else content)

    def end(self, name, newline=True):
        self._write('</%s>%s' % (name, self._line_separator if newline else ''))

    def element(self, name, content=None, attrs=None, escape=True,
                newline=True):
        self.start(name, attrs, newline=False)
        self.content(content, escape)
        self.end(name, newline)

    def close(self):
        self.output.close()

    def _write(self, text):
        self.output.write(self._encode(text))

    def _get_attrs(self, attrs):
        if not attrs:
            return ''
        return ' ' + ' '.join(self._format_attributes(attrs))

class HtmlWriter(_MarkupWriter):

    def _escape(self, content):
        return html_escape(content)

    def _format_attributes(self, attrs):
        return ('%s="%s"' % (name, attribute_escape(attrs[name]))
                             for name in sorted(attrs))

class XmlWriter(_MarkupWriter):

    def __init__(self, output, line_separator=os.linesep, encoding=None):
        _MarkupWriter.__init__(self, output, line_separator, encoding)
        self._preamble()

    def _preamble(self):
        self.content('<?xml version="1.0" encoding="UTF-8"?>\n', escape=False)

    def _escape(self, text):
        return xml_escape(text)

    def _format_attributes(self, attrs):
        return ('%s="%s"' % (name, attribute_escape(unicode(attrs[name])))
                             for name in attrs)
<commit_msg>markupwriter: api documentation<commit_after>#  Copyright 2008-2012 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import os

from .markuputils import html_escape, xml_escape, attribute_escape

class _MarkupWriter(object):

    def __init__(self, output, line_separator=os.linesep, encoding=None):
        """Creates new _MarkupWriter.

        :param output: Either an opened, file like object, or a path to the
            desired output file. In the latter case, the file is created.
        :param line_separator: Defines the used line separator.
        :param encoding: Encoding to be used to encode all text written to the
            output file. If `None`, text will not be encoded.
        """
        if isinstance(output, basestring):
            output = open(output, 'w')
        self.output = output
        self._line_separator = line_separator
        self._encode = self._create_encoder(encoding)

    def _create_encoder(self, encoding):
        if encoding is None:
            return lambda text: text
        return lambda text: text.encode(encoding)

    def start(self, name, attrs=None, newline=True):
        self._write('<%s%s>%s' % (name, self._get_attrs(attrs),
                                  self._line_separator if newline else ''))

    def content(self, content=None, escape=True):
        if content:
            self._write(self._escape(content) if escape else content)

    def end(self, name, newline=True):
        self._write('</%s>%s' % (name, self._line_separator if newline else ''))

    def element(self, name, content=None, attrs=None, escape=True,
                newline=True):
        self.start(name, attrs, newline=False)
        self.content(content, escape)
        self.end(name, newline)

    def close(self):
        self.output.close()

    def _write(self, text):
        self.output.write(self._encode(text))

    def _get_attrs(self, attrs):
        if not attrs:
            return ''
        return ' ' + ' '.join(self._format_attributes(attrs))

class HtmlWriter(_MarkupWriter):

    def _escape(self, content):
        return html_escape(content)

    def _format_attributes(self, attrs):
        return ('%s="%s"' % (name, attribute_escape(attrs[name]))
                             for name in sorted(attrs))

class XmlWriter(_MarkupWriter):

    def __init__(self, output, line_separator=os.linesep, encoding=None):
        _MarkupWriter.__init__(self, output, line_separator, encoding)
        self._preamble()

    def _preamble(self):
        self.content('<?xml version="1.0" encoding="UTF-8"?>\n', escape=False)

    def _escape(self, text):
        return xml_escape(text)

    def _format_attributes(self, attrs):
        return ('%s="%s"' % (name, attribute_escape(unicode(attrs[name])))
                             for name in attrs)
<|endoftext|>


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

""" OneLogin_Saml2_Authn_Request class

Copyright (c) 2014, OneLogin, Inc.
All rights reserved.

AuthNRequest class of OneLogin's Python Toolkit.

"""
from base64 import b64encode

from onelogin.saml2.constants import OneLogin_Saml2_Constants
from onelogin.saml2.utils import OneLogin_Saml2_Utils

class OneLogin_Saml2_Authn_Request(object):
    """

    This class handles an AuthNRequest. It builds an
    AuthNRequest object.

    """

    def __init__(self, settings, force_authn=False, is_passive=False):
        """
        Constructs the AuthnRequest object.

        :param settings: OSetting data
        :type return_to: OneLogin_Saml2_Settings

        :param force_authn: Optional argument. When true the AuthNReuqest will set the ForceAuthn='true'.
        :type force_authn: bool

        :param is_passive: Optional argument. When true the AuthNReuqest will set the Ispassive='true'.
        :type is_passive: bool
        """
        self.__settings = settings

        sp_data = self.__settings.get_sp_data()
        idp_data = self.__settings.get_idp_data()
        security = self.__settings.get_security_data()

        uid = OneLogin_Saml2_Utils.generate_unique_id()
        self.__id = uid
        issue_instant = OneLogin_Saml2_Utils.parse_time_to_SAML(OneLogin_Saml2_Utils.now())

        destination = idp_data['singleSignOnService']['url']

        name_id_policy_format = None
        if 'wantNameId' in security and security['wantNameId']:
            name_id_policy_format = sp_data['NameIDFormat']
            if 'wantNameIdEncrypted' in security and security['wantNameIdEncrypted']:
                name_id_policy_format = OneLogin_Saml2_Constants.NAMEID_ENCRYPTED

        name_id_policy = ''
        if name_id_policy_format:
            name_id_policy = """<samlp:NameIDPolicy Format="%(name_id_policy)s" AllowCreate="true" />""" % name_id_policy_format

        provider_name_str = ''
        organization_data = settings.get_organization()
        if isinstance(organization_data, dict) and organization_data:
            langs = organization_data.keys()
            if 'en-US' in langs:
                lang = 'en-US'
            else:
                lang = langs[0]
            if 'displayname' in organization_data[lang] and organization_data[lang]['displayname'] is not None:
                provider_name_str = 'ProviderName="%s"' % organization_data[lang]['displayname']

        force_authn_str = ''
        if force_authn is True:
            force_authn_str = 'ForceAuthn="true"'

        is_passive_str = ''
        if is_passive is True:
            is_passive_str = 'IsPassive="true"'

        requested_authn_context_str = ''
        if 'requestedAuthnContext' in security.keys() and security['requestedAuthnContext'] is not False:
            authn_comparison = 'exact'
            if 'requestedAuthnContextComparison' in security.keys():
                authn_comparison = security['requestedAuthnContextComparison']

            if security['requestedAuthnContext'] is True:
                if 'activeDirectoryFederation' in security and security['activeDirectoryFederation']:
                    authn_context = """<saml:AuthnContextClassRef xmlns:samlp="urn:oasis:names:tc:SAML:2.0:assertion">urn:federation:authentication:windows</saml:AuthnContextClassRef>"""
                else:
                    authn_context = """<saml:AuthnContextClassRef>urn:oasis:names:tc:SAML:2.0:ac:classes:PasswordProtectedTransport</saml:AuthnContextClassRef>"""
                requested_authn_context_str = """    <samlp:RequestedAuthnContext Comparison="%s">%s</samlp:RequestedAuthnContext>""" % (authn_comparison, authn_context)
            else:
                requested_authn_context_str = '     <samlp:RequestedAuthnContext Comparison="%s">' % authn_comparison
                for authn_context in security['requestedAuthnContext']:
                    requested_authn_context_str += '<saml:AuthnContextClassRef>%s</saml:AuthnContextClassRef>' % authn_context
                requested_authn_context_str += '    </samlp:RequestedAuthnContext>'

        attr_consuming_service_str = ''
        if 'attributeConsumingService' in sp_data and sp_data['attributeConsumingService']:
            attr_consuming_service_str = 'AttributeConsumingServiceIndex="1"'

        request = """<samlp:AuthnRequest
    xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol"
    xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion"
    ID="%(id)s"
    Version="2.0"
    %(provider_name)s
    %(force_authn_str)s
    %(is_passive_str)s
    IssueInstant="%(issue_instant)s"
    Destination="%(destination)s"
    ProtocolBinding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST"
    AssertionConsumerServiceURL="%(assertion_url)s"
    %(attr_consuming_service_str)s>
    <saml:Issuer>%(entity_id)s</saml:Issuer> %(name_id_policy)s
%(requested_authn_context_str)s
</samlp:AuthnRequest>""" % \
            {
                'id': uid,
                'provider_name': provider_name_str,
                'force_authn_str': force_authn_str,
                'is_passive_str': is_passive_str,
                'issue_instant': issue_instant,
                'destination': destination,
                'assertion_url': sp_data['assertionConsumerService']['url'],
                'entity_id': sp_data['entityId'],
                'name_id_policy': name_id_policy,
                'requested_authn_context_str': requested_authn_context_str,
                'attr_consuming_service_str': attr_consuming_service_str
            }

        self.__authn_request = request

    def get_request(self, deflate=True):
        """
        Returns unsigned AuthnRequest.
        :param deflate: It makes the deflate process optional
        :type: bool
        :return: AuthnRequest maybe deflated and base64 encoded
        :rtype: str object
        """
        if deflate:
            request = OneLogin_Saml2_Utils.deflate_and_base64_encode(self.__authn_request)
        else:
            request = b64encode(self.__authn_request)
        return request

    def get_id(self):
        """
        Returns the AuthNRequest ID.
        :return: AuthNRequest ID
        :rtype: string
        """
        return self.__id
<commit_msg>Fixes to ensure all tests pass<commit_after># -*- coding: utf-8 -*-

""" OneLogin_Saml2_Authn_Request class

Copyright (c) 2014, OneLogin, Inc.
All rights reserved.

AuthNRequest class of OneLogin's Python Toolkit.

"""
from base64 import b64encode

from onelogin.saml2.constants import OneLogin_Saml2_Constants
from onelogin.saml2.utils import OneLogin_Saml2_Utils

class OneLogin_Saml2_Authn_Request(object):
    """

    This class handles an AuthNRequest. It builds an
    AuthNRequest object.

    """

    def __init__(self, settings, force_authn=False, is_passive=False):
        """
        Constructs the AuthnRequest object.

        :param settings: OSetting data
        :type return_to: OneLogin_Saml2_Settings

        :param force_authn: Optional argument. When true the AuthNReuqest will set the ForceAuthn='true'.
        :type force_authn: bool

        :param is_passive: Optional argument. When true the AuthNReuqest will set the Ispassive='true'.
        :type is_passive: bool
        """
        self.__settings = settings

        sp_data = self.__settings.get_sp_data()
        idp_data = self.__settings.get_idp_data()
        security = self.__settings.get_security_data()

        uid = OneLogin_Saml2_Utils.generate_unique_id()
        self.__id = uid
        issue_instant = OneLogin_Saml2_Utils.parse_time_to_SAML(OneLogin_Saml2_Utils.now())

        destination = idp_data['singleSignOnService']['url']

        name_id_policy_format = sp_data['NameIDFormat']
        if 'wantNameIdEncrypted' in security and security['wantNameIdEncrypted']:
            name_id_policy_format = OneLogin_Saml2_Constants.NAMEID_ENCRYPTED

        name_id_policy = ''
        if security.get('wantNameId', True):
            name_id_policy = """<samlp:NameIDPolicy Format="%s" AllowCreate="true" />""" % name_id_policy_format

        provider_name_str = ''
        organization_data = settings.get_organization()
        if isinstance(organization_data, dict) and organization_data:
            langs = organization_data.keys()
            if 'en-US' in langs:
                lang = 'en-US'
            else:
                lang = langs[0]
            if 'displayname' in organization_data[lang] and organization_data[lang]['displayname'] is not None:
                provider_name_str = 'ProviderName="%s"' % organization_data[lang]['displayname']

        force_authn_str = ''
        if force_authn is True:
            force_authn_str = 'ForceAuthn="true"'

        is_passive_str = ''
        if is_passive is True:
            is_passive_str = 'IsPassive="true"'

        requested_authn_context_str = ''
        if 'requestedAuthnContext' in security.keys() and security['requestedAuthnContext'] is not False:
            authn_comparison = 'exact'
            if 'requestedAuthnContextComparison' in security.keys():
                authn_comparison = security['requestedAuthnContextComparison']

            if security['requestedAuthnContext'] is True:
                if 'activeDirectoryFederation' in security and security['activeDirectoryFederation']:
                    authn_context = """<saml:AuthnContextClassRef xmlns:samlp="urn:oasis:names:tc:SAML:2.0:assertion">urn:federation:authentication:windows</saml:AuthnContextClassRef>"""
                else:
                    authn_context = """<saml:AuthnContextClassRef>urn:oasis:names:tc:SAML:2.0:ac:classes:PasswordProtectedTransport</saml:AuthnContextClassRef>"""
                requested_authn_context_str = """    <samlp:RequestedAuthnContext Comparison="%s">%s</samlp:RequestedAuthnContext>""" % (authn_comparison, authn_context)
            else:
                requested_authn_context_str = '     <samlp:RequestedAuthnContext Comparison="%s">' % authn_comparison
                for authn_context in security['requestedAuthnContext']:
                    requested_authn_context_str += '<saml:AuthnContextClassRef>%s</saml:AuthnContextClassRef>' % authn_context
                requested_authn_context_str += '    </samlp:RequestedAuthnContext>'

        attr_consuming_service_str = ''
        if 'attributeConsumingService' in sp_data and sp_data['attributeConsumingService']:
            attr_consuming_service_str = 'AttributeConsumingServiceIndex="1"'

        request = """<samlp:AuthnRequest
    xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol"
    xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion"
    ID="%(id)s"
    Version="2.0"
    %(provider_name)s
    %(force_authn_str)s
    %(is_passive_str)s
    IssueInstant="%(issue_instant)s"
    Destination="%(destination)s"
    ProtocolBinding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST"
    AssertionConsumerServiceURL="%(assertion_url)s"
    %(attr_consuming_service_str)s>
    <saml:Issuer>%(entity_id)s</saml:Issuer> %(name_id_policy)s
%(requested_authn_context_str)s
</samlp:AuthnRequest>""" % \
            {
                'id': uid,
                'provider_name': provider_name_str,
                'force_authn_str': force_authn_str,
                'is_passive_str': is_passive_str,
                'issue_instant': issue_instant,
                'destination': destination,
                'assertion_url': sp_data['assertionConsumerService']['url'],
                'entity_id': sp_data['entityId'],
                'name_id_policy': name_id_policy,
                'requested_authn_context_str': requested_authn_context_str,
                'attr_consuming_service_str': attr_consuming_service_str
            }

        self.__authn_request = request

    def get_request(self, deflate=True):
        """
        Returns unsigned AuthnRequest.
        :param deflate: It makes the deflate process optional
        :type: bool
        :return: AuthnRequest maybe deflated and base64 encoded
        :rtype: str object
        """
        if deflate:
            request = OneLogin_Saml2_Utils.deflate_and_base64_encode(self.__authn_request)
        else:
            request = b64encode(self.__authn_request)
        return request

    def get_id(self):
        """
        Returns the AuthNRequest ID.
        :return: AuthNRequest ID
        :rtype: string
        """
        return self.__id
<|endoftext|>


<commit_before>"""Module for reading Translational and rotational data from a
    file to compute dynamics quantities.
"""

import glob
from CompDynamics import CompRotDynamics
from TransData import TransRotData

def compute_file(fname, outfile='out.dat'):
    """Read data from file, compute dynamics and print results

    The file is read line by line, where the data from each frame should
    reside on a single line. While not sticking completely to the JSON format
    it allows for file sizes larger than would fit into memory.

    Note:
        This assumes that the data being read from a file is a
        class:`TransData.TransRotData` object.

    Args:
        fname (string): Path to file containing the translations and rotations
        outfile (string): Filename of output file
    """
    with open(fname) as infile:
        init = True
        for line in infile:
            data = TransRotData()
            data.from_json(line)
            dyn = CompRotDynamics(data)
            if init:
                dyn.print_heading(outfile)
                init = False
            dyn.print_all(outfile)

def compute_all(pattern="*-tr.dat", suffix="-dyn.dat", directory="."):
    """Compute all files in directory matching pattern

    This finds all the files in a specified directory that match a pattern,
    computing all the dynamic quanities on the translational and rotational
    data in that file.

    Todo:
        * Spawn the analysis into multiple threads

    Args:
        pattern (str): Pattern that matches the files to compute. Defaults to
            `"*-tr.dat"`.
        suffix (str): replace the last characters of the file with this.
            Default is `"-dyn.dat"`.
        directory (str): The directory in which to search for the files.
            Default is the current directory.
    """
    files = glob.glob(directory+"/"+pattern)
    for infile in files:
        compute_file(infile, infile[:-(len(pattern)-1)]+suffix)

if __name__ == "__main__":
    compute_all()
<commit_msg>File progress in  FileDynamics module<commit_after>"""Module for reading Translational and rotational data from a
    file to compute dynamics quantities.
"""

import glob
from CompDynamics import CompRotDynamics
from TransData import TransRotData

def compute_file(fname, outfile='out.dat'):
    """Read data from file, compute dynamics and print results

    The file is read line by line, where the data from each frame should
    reside on a single line. While not sticking completely to the JSON format
    it allows for file sizes larger than would fit into memory.

    Note:
        This assumes that the data being read from a file is a
        class:`TransData.TransRotData` object.

    Args:
        fname (string): Path to file containing the translations and rotations
        outfile (string): Filename of output file
    """
    with open(fname) as infile:
        init = True
        for line in infile:
            data = TransRotData()
            data.from_json(line)
            dyn = CompRotDynamics(data)
            if init:
                dyn.print_heading(outfile)
                init = False
            dyn.print_all(outfile)

def compute_all(pattern="*-tr.dat", suffix="-dyn.dat", directory="."):
    """Compute all files in directory matching pattern

    This finds all the files in a specified directory that match a pattern,
    computing all the dynamic quanities on the translational and rotational
    data in that file.

    Todo:
        * Spawn the analysis into multiple threads

    Args:
        pattern (str): Pattern that matches the files to compute. Defaults to
            `"*-tr.dat"`.
        suffix (str): replace the last characters of the file with this.
            Default is `"-dyn.dat"`.
        directory (str): The directory in which to search for the files.
            Default is the current directory.
    """
    files = glob.glob(directory+"/"+pattern)
    for infile in files:
        print(infile)
        compute_file(infile, infile[:-(len(pattern)-1)]+suffix)

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


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

# This script retrieves data about the use of classes and properties
# on Wikidata from SPARQL. The results are processed and stored in the
# files properties.json and classes.json. The script must be run in a
# directory that already contains these files (possibly with no content
# other than an empty map {}). Normally, the files are first generated
# by the Java export tool, which provides some statistics that cannot
# be obtained in an acceptable time from SPARQL.

from __future__ import print_function

import requests
import json
import os
import time
#import pprint

SPARQL_SERVICE_URL = 'https://query.wikidata.org/sparql'

def doSparqlQuery(query):
	r = requests.get(SPARQL_SERVICE_URL, params={'query': "#TOOL:SQID Python Helper\n" + query, 'format': 'json'});
	return json.loads(r.text)

def sparqlQuery(query):
	try:
		return doSparqlQuery(query)
	except ValueError:
		print("SPARQL query failed, possibly because of a time out. Waiting for 60 sec ...\n")
		time.sleep(60)
		print("Retrying query ...\n")
		try:
			return doSparqlQuery(query)
		except ValueError:
			print("Failed to get SPARQL results. Giving up.")
			return json.loads('{ "error": "timeout", "results": { "bindings": {} } }')

def setPropertyStatistics(propertyStatistics, propertyId, statisticType, numberString):
	try:
		if propertyId in propertyStatistics:
			propertyStatistics[propertyId][statisticType] = int(numberString)
	except ValueError:
		print("Error reading count value " + numberString + ": not an integer number")

def storeStatistics(key, value):
	with open('statistics.json', 'r') as outfile:
		data = json.load(outfile)
		data[key] = value
		with open('statistics.json', 'w') as outfile:
			json.dump(data, outfile)

def updateClassRecords() :
	startTime = time.strftime("%Y-%m-%dT%H:%M:%S")
	print("[" + startTime + "] Fetching class ids and labels for classes with direct instances ...")

	resultsClasses = sparqlQuery("""SELECT ?cl ?clLabel ?c
WHERE {
	{ SELECT ?cl (count(*) as ?c) WHERE { ?i wdt:P31 ?cl } GROUP BY ?cl }
	SERVICE wikibase:label {
			bd:serviceParam wikibase:language "en" .
	}
}""")

	if ( "error" in resultsClasses ):
		print("Failed to fetch class data from SPARQL. Trying simpler query ...")
		resultsClasses = sparqlQuery("""SELECT ?cl ?clLabel
WHERE {
	{ SELECT DISTINCT ?cl WHERE { ?i wdt:P31 ?cl } }
	SERVICE wikibase:label {
			bd:serviceParam wikibase:language "en" .
	}
}""")
		if ( "error" in resultsClasses ):
			print("Failed to fetch class data from SPARQL. Giving up.")
			return

	print("Augmenting current class data ...")
	classesWithInstances = {}
	with open('classes.json') as classFile:
		data = json.load(classFile)
		for binding in resultsClasses['results']['bindings']:
			classUri = binding['cl']['value']
			if classUri[0:31] == 'http://www.wikidata.org/entity/':
				classId = classUri[32:]
				classesWithInstances[classId] = True
				classLabel = binding['clLabel']['value']
				if classId in data:
					data[classId]['l'] = classLabel
					if "c" in binding:
						data[classId]['i'] = int(binding['c']['value'])
				else:
					if "c" in binding:
						data[classId] = {'l':classLabel, 'i':int(binding['c']['value'])}
					else:
						data[classId] = {'l':classLabel}
		# zero instance counts for classes not found in the query:
		for classId in data :
			if classId not in classesWithInstances :
				data[classId]['i'] = 0

		print("Writing json class data ...")
		with open('classes-new.json', 'w') as outfile:
			json.dump(data, outfile,separators=(',', ':'))

		print("Replacing classes json file ...")
		os.rename("classes-new.json","classes.json")

	storeStatistics('classUpdate',startTime)
	print("Class update complete.")

def updatePropertyRecords() :
	startTime = time.strftime("%Y-%m-%dT%H:%M:%S")
	print("[" + startTime + "] Fetching property ids, labels, and types ...")

	resultsProperties = sparqlQuery("""PREFIX wikibase: <http://wikiba.se/ontology#>
	SELECT ?id ?idLabel ?type
	WHERE {
		?id a wikibase:Property .
		?id wikibase:propertyType ?type
		SERVICE wikibase:label {
			bd:serviceParam wikibase:language "en" .
		}
	}""")

	propertyStatistics = {}
	for binding in resultsProperties['results']['bindings']:
		propertyUri = binding['id']['value']
		if propertyUri[0:31] == 'http://www.wikidata.org/entity/':
			propertyLabel = binding['idLabel']['value']
			typeUri = binding['type']['value']
			if typeUri[0:26] == 'http://wikiba.se/ontology#':
				typeName = typeUri[26:]
			else:
				typeName = 'UNKNOWN'
			propertyStatistics[propertyUri[32:]] = {'l': propertyLabel, 't': typeName, 's': 0, 'q': 0, 'r': 0}
		else:
			print('Error reading property URI ' + propertyUri)

	print("Fetching property usage statistics ...")

	resultsPropertyCounts = sparqlQuery("""SELECT ?p (count(*) as ?c)
	WHERE {
		?s ?p ?o .
	}
	GROUP BY ?p
	ORDER BY DESC(?c)""");

	for binding in resultsPropertyCounts['results']['bindings']:
		propertyUri = binding['p']['value']
		# Note: Always check for longest matches first
		# Checking for reference/value and quantifier/value is not necessary to get overall counts
		#if len(propertyUri)>45 and propertyUri[0:45] == 'http://www.wikidata.org/prop/reference/value/':
			#propertyId = propertyUri[46:]
			#setPropertyStatistics(propertyStatistics, propertyId,'rv', binding['c']['value'])
		#elif len(propertyUri)>45 and propertyUri[0:45] == 'http://www.wikidata.org/prop/qualifier/value/':
			#propertyId = propertyUri[46:]
			#setPropertyStatistics(propertyStatistics, propertyId,'qv', binding['c']['value'])
		if len(propertyUri)>39 and propertyUri[0:39] == 'http://www.wikidata.org/prop/reference/':
			propertyId = propertyUri[40:]
			setPropertyStatistics(propertyStatistics, propertyId,'r', binding['c']['value'])
		elif len(propertyUri)>39 and propertyUri[0:39] == 'http://www.wikidata.org/prop/qualifier/':
			propertyId = propertyUri[40:]
			setPropertyStatistics(propertyStatistics, propertyId,'q', binding['c']['value'])
		elif len(propertyUri)>29 and propertyUri[0:29] == 'http://www.wikidata.org/prop/':
			propertyId = propertyUri[30:]
			setPropertyStatistics(propertyStatistics, propertyId,'s', binding['c']['value'])
		#else:
		#	# The rest are Wikibase ontology properties and some RDF(S) and schema.org URIs:
		#	print binding['p']['value'] + ' = ' + binding['c']['value']

	print("Augmenting current property data ...")
	with open('properties.json') as propertyFile:
		data = json.load(propertyFile)

		for propertyId, propertyData in propertyStatistics.items():
			if propertyId in data:
				data[propertyId]['l'] = propertyData['l']
				data[propertyId]['d'] = propertyData['t']
				data[propertyId]['s'] = propertyData['s']
				data[propertyId]['q'] = propertyData['q']
				data[propertyId]['e'] = propertyData['r']
			else:
				data[propertyId] = { 'l': propertyData['l'], 'd': propertyData['t'], 's': propertyData['s'], 'q': propertyData['q'], 'e': propertyData['r']}

		print("Writing json property data ...")
		with open('properties-new.json', 'w') as outfile:
			json.dump(data, outfile,separators=(',', ':'))

		print("Replacing properties json file ...")
		os.rename("properties-new.json","properties.json")

	storeStatistics('propertyUpdate',startTime)
	print("Property update complete.")

#pprint.pprint(propertyStatistics)
<commit_msg>Behave gracefully if run without any data<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

# This script retrieves data about the use of classes and properties
# on Wikidata from SPARQL. The results are processed and stored in the
# files properties.json and classes.json. The script must be run in a
# directory that already contains these files (possibly with no content
# other than an empty map {}). Normally, the files are first generated
# by the Java export tool, which provides some statistics that cannot
# be obtained in an acceptable time from SPARQL.

from __future__ import print_function

import requests
import json
import os
import time
#import pprint

SPARQL_SERVICE_URL = 'https://query.wikidata.org/sparql'

def doSparqlQuery(query):
	r = requests.get(SPARQL_SERVICE_URL, params={'query': "#TOOL:SQID Python Helper\n" + query, 'format': 'json'});
	return json.loads(r.text)

def sparqlQuery(query):
	try:
		return doSparqlQuery(query)
	except ValueError:
		print("SPARQL query failed, possibly because of a time out. Waiting for 60 sec ...\n")
		time.sleep(60)
		print("Retrying query ...\n")
		try:
			return doSparqlQuery(query)
		except ValueError:
			print("Failed to get SPARQL results. Giving up.")
			return json.loads('{ "error": "timeout", "results": { "bindings": {} } }')

def setPropertyStatistics(propertyStatistics, propertyId, statisticType, numberString):
	try:
		if propertyId in propertyStatistics:
			propertyStatistics[propertyId][statisticType] = int(numberString)
	except ValueError:
		print("Error reading count value " + numberString + ": not an integer number")

def storeStatistics(key, value):
	try:
		with open('statistics.json', 'r') as outfile:
			data = json.load(outfile)
	except OSError:
		data = {}

	data[key] = value
	with open('statistics.json', 'w') as outfile:
		json.dump(data, outfile)

def updateClassRecords() :
	startTime = time.strftime("%Y-%m-%dT%H:%M:%S")
	print("[" + startTime + "] Fetching class ids and labels for classes with direct instances ...")

	resultsClasses = sparqlQuery("""SELECT ?cl ?clLabel ?c
WHERE {
	{ SELECT ?cl (count(*) as ?c) WHERE { ?i wdt:P31 ?cl } GROUP BY ?cl }
	SERVICE wikibase:label {
			bd:serviceParam wikibase:language "en" .
	}
}""")

	if ( "error" in resultsClasses ):
		print("Failed to fetch class data from SPARQL. Trying simpler query ...")
		resultsClasses = sparqlQuery("""SELECT ?cl ?clLabel
WHERE {
	{ SELECT DISTINCT ?cl WHERE { ?i wdt:P31 ?cl } }
	SERVICE wikibase:label {
			bd:serviceParam wikibase:language "en" .
	}
}""")
		if ( "error" in resultsClasses ):
			print("Failed to fetch class data from SPARQL. Giving up.")
			return

	print("Augmenting current class data ...")
	classesWithInstances = {}
	try:
		with open('classes.json') as classFile:
			data = json.load(classFile)
	except OSError:
		data = {}

	for binding in resultsClasses['results']['bindings']:
		classUri = binding['cl']['value']
		if classUri[0:31] == 'http://www.wikidata.org/entity/':
			classId = classUri[32:]
			classesWithInstances[classId] = True
			classLabel = binding['clLabel']['value']
			if classId in data:
				data[classId]['l'] = classLabel
				if "c" in binding:
					data[classId]['i'] = int(binding['c']['value'])
			else:
				if "c" in binding:
					data[classId] = {'l':classLabel, 'i':int(binding['c']['value'])}
				else:
					data[classId] = {'l':classLabel}
	# zero instance counts for classes not found in the query:
	for classId in data:
		if classId not in classesWithInstances:
			data[classId]['i'] = 0

	print("Writing json class data ...")
	with open('classes-new.json', 'w') as outfile:
		json.dump(data, outfile,separators=(',', ':'))

	print("Replacing classes json file ...")
	os.rename("classes-new.json","classes.json")

	storeStatistics('classUpdate',startTime)
	print("Class update complete.")

def updatePropertyRecords() :
	startTime = time.strftime("%Y-%m-%dT%H:%M:%S")
	print("[" + startTime + "] Fetching property ids, labels, and types ...")

	resultsProperties = sparqlQuery("""PREFIX wikibase: <http://wikiba.se/ontology#>
	SELECT ?id ?idLabel ?type
	WHERE {
		?id a wikibase:Property .
		?id wikibase:propertyType ?type
		SERVICE wikibase:label {
			bd:serviceParam wikibase:language "en" .
		}
	}""")

	propertyStatistics = {}
	for binding in resultsProperties['results']['bindings']:
		propertyUri = binding['id']['value']
		if propertyUri[0:31] == 'http://www.wikidata.org/entity/':
			propertyLabel = binding['idLabel']['value']
			typeUri = binding['type']['value']
			if typeUri[0:26] == 'http://wikiba.se/ontology#':
				typeName = typeUri[26:]
			else:
				typeName = 'UNKNOWN'
			propertyStatistics[propertyUri[32:]] = {'l': propertyLabel, 't': typeName, 's': 0, 'q': 0, 'r': 0}
		else:
			print('Error reading property URI ' + propertyUri)

	print("Fetching property usage statistics ...")

	resultsPropertyCounts = sparqlQuery("""SELECT ?p (count(*) as ?c)
	WHERE {
		?s ?p ?o .
	}
	GROUP BY ?p
	ORDER BY DESC(?c)""");

	for binding in resultsPropertyCounts['results']['bindings']:
		propertyUri = binding['p']['value']
		# Note: Always check for longest matches first
		# Checking for reference/value and quantifier/value is not necessary to get overall counts
		#if len(propertyUri)>45 and propertyUri[0:45] == 'http://www.wikidata.org/prop/reference/value/':
			#propertyId = propertyUri[46:]
			#setPropertyStatistics(propertyStatistics, propertyId,'rv', binding['c']['value'])
		#elif len(propertyUri)>45 and propertyUri[0:45] == 'http://www.wikidata.org/prop/qualifier/value/':
			#propertyId = propertyUri[46:]
			#setPropertyStatistics(propertyStatistics, propertyId,'qv', binding['c']['value'])
		if len(propertyUri)>39 and propertyUri[0:39] == 'http://www.wikidata.org/prop/reference/':
			propertyId = propertyUri[40:]
			setPropertyStatistics(propertyStatistics, propertyId,'r', binding['c']['value'])
		elif len(propertyUri)>39 and propertyUri[0:39] == 'http://www.wikidata.org/prop/qualifier/':
			propertyId = propertyUri[40:]
			setPropertyStatistics(propertyStatistics, propertyId,'q', binding['c']['value'])
		elif len(propertyUri)>29 and propertyUri[0:29] == 'http://www.wikidata.org/prop/':
			propertyId = propertyUri[30:]
			setPropertyStatistics(propertyStatistics, propertyId,'s', binding['c']['value'])
		#else:
		#	# The rest are Wikibase ontology properties and some RDF(S) and schema.org URIs:
		#	print binding['p']['value'] + ' = ' + binding['c']['value']

	print("Augmenting current property data ...")
	try:
		with open('properties.json') as propertyFile:
			data = json.load(propertyFile)
	except OSError:
		data = {}

	for propertyId, propertyData in propertyStatistics.items():
		if propertyId in data:
			data[propertyId]['l'] = propertyData['l']
			data[propertyId]['d'] = propertyData['t']
			data[propertyId]['s'] = propertyData['s']
			data[propertyId]['q'] = propertyData['q']
			data[propertyId]['e'] = propertyData['r']
		else:
			data[propertyId] = { 'l': propertyData['l'], 'd': propertyData['t'], 's': propertyData['s'], 'q': propertyData['q'], 'e': propertyData['r']}

	print("Writing json property data ...")
	with open('properties-new.json', 'w') as outfile:
		json.dump(data, outfile,separators=(',', ':'))

	print("Replacing properties json file ...")
	os.rename("properties-new.json","properties.json")

	storeStatistics('propertyUpdate',startTime)
	print("Property update complete.")

#pprint.pprint(propertyStatistics)
<|endoftext|>


<commit_before>import inspect
import sys

import pandas as pd

pdDatetimeIndex = pd.tseries.index.DatetimeIndex
pdInt64Index = pd.core.index.Int64Index
pdCoreIndex = pd.core.index.Index

import sqlalchemy as sqla

import datetime as dt

class IndexImplementer(object):

    """
    IndexImplementer is the base required to implement
    an index of a specific type.  The
    same instance is created at two points in
    the Trump dataflow:

    1. the datatable getting cached and
    2. the data being served.

    The IndexImplementer should be
    indempotent, and dataframe/series agnostic.
    """

    sqlatyp = sqla.Integer
    pytyp = int
    pindt = pd.Index

    def __init__(self, case, **kwargs):
        """
        :param case: str
            This should match a case used to switch
            the logic created in each subclass of IndexImplementer
        :param kwargs: dict
        """

        self.case = case
        self.k = kwargs
    def orfs_ind_from_str(self, userinput):
        ui = {}
        exec("ui = " + userinput)
        obj = self.pytyp(**ui)
        return obj
    def create_empty(self):
        return self.pindt([])
    def build_ordf(self, orind, orval, colname):
        ordf = pd.DataFrame(index=orind, data=orval, columns=[colname])
        return self._common_passthrough(ordf)
    def _common_passthrough(self, obj):
        try:
            return getattr(self, '_' + self.case)(obj)
        except AttributeError:
            self._notimp_error()
    def _notimp_error(self):
        msg = "Indexing case '{}' unsupported for Index Implementer {}"
        msg = msg.format(self.case, self.__class__.__name__)
        raise NotImplementedError(msg)
    def _notimpobj_error(self, obj):
        msg = "Object of type '{}' unsupported for Index Implementer {}, case {}"
        msg = msg.format(type(obj), self.__class__.__name__, self.case)
        raise TypeError(msg)

    def process_post_db(self, df):
        return self._common_passthrough(df)
    def process_post_feed_cache(self, s):
        return self._common_passthrough(s)
    def process_post_orfs(self, df):
        return self._common_passthrough(df)

class DateIndexImp(IndexImplementer):
    sqlatyp = sqla.Date
    pytyp = dt.date
    pindt = pd.Index
    def __init__(self, dfors, case, kwargs):
        raise NotImplemented()

class DatetimeIndexImp(IndexImplementer):
    """
    Implements a pandas DatetimeIndex

    Cases include:

    * **asis** - Cache timestamps to the database and drop
      any intelligence associated with the index, such as frequency.
      serve a Series with a DatetimeIndex, without frequency.

      If the index consists of 4-digit integers, it will
      be treated as the year, in a date which is of the form
      YYYY-12-31.

    * **asfreq** - Apply 'asfreq' logic prior to cache, and
      apply the same logic when serving.

    * **date_range** - Create a new index, using pandas date_range(), at
      time of cache...  NotImplemented yet.

    * **guess** - NotImplemented yet.
      Attempt to guess the frequency at time of cache,
      and time of serve.

    * **guess_post** - NotImplemented yet.
      Attempt to guess the frequency at time of serve,
      but store the cache unsaved.

    In the event that case hasn't implemented the logic
    to handle a specific datatype, a rudimentary
    attempt to convert it to a DatetimeIndex is applied
    by inspecting the start and end, with the kwargs. passed
    pandas.DatetimeIndex constructor.

    """
    sqlatyp = sqla.DateTime
    pytyp = dt.datetime
    pindt = pd.DatetimeIndex

    def _asis(self, obj):
        if isinstance(obj.index, pdDatetimeIndex):
            return obj
        elif isinstance(obj.index, pdInt64Index):
            if all(len(str(i)) == 4 for i in obj.index):
                #safe to assume we meant years...
                newind = [dt.date(y, 12, 31) for y in obj.index]
                obj.index = newind
                obj.index = obj.index.to_datetime()
                return obj
            self._notimpobj_error(obj)
        elif isinstance(obj.index, pdCoreIndex):
            if isinstance(obj.index[0], dt.date):
                newind = pdDatetimeIndex(obj.index)
                obj.index = newind
                return obj
            self._notimpobj_error(obj)
        else:
            return self._default(obj, **self.k)

    def _asfreq(self, obj):
        if isinstance(obj.index, pdDatetimeIndex):
            obj = obj.asfreq(**self.k)
            return obj
        elif isinstance(obj.index, pdInt64Index):
            if all(len(str(i)) == 4 for i in obj.index):
                #safe to assume we meant years...
                newind = [dt.date(y, 12, 31) for y in obj.index]
                obj.index = newind
                obj.index = obj.index.to_datetime()
                obj = obj.asfreq(**self.k)
                return obj
            self._notimpobj_error(obj)
        elif isinstance(obj.index[0], (str, unicode)):
            newind = pd.DatetimeIndex(obj.index)
            obj.index = newind
            obj = obj.asfreq(**self.k)
            return obj
        else:
            return self._default(obj, **self.k)

    @staticmethod
    def _default(obj, **kwargs):
        start = pd.to_datetime(obj.index[0])
        end = pd.to_datetime(obj.index[-1])
        newind = pdDatetimeIndex(start=start, end=end, **kwargs)
        obj = obj.reindex(newind)
        return obj

class PeriodIndexImp(IndexImplementer):
    """
    Implements a pandas PeriodIndex

    NotImplemented, yet.
    """
    sqlatyp = sqla.DateTime
    pytyp = pd.Period
    pindt = pd.PeriodIndex

    def __init__(self, dfors, case, kwargs):
        raise NotImplementedError()

class IntIndexImp(IndexImplementer):
    """
    Implements a pandas Int64Index.

    Cases include:

    * **asis** - attempts to pass the index through,
      without applying any logic.  Use this, if the index
      is already integers, or unique and integer-like.

    * **drop** - will drop the pandas index, to reset it.

    """

    sqlatyp = sqla.Integer
    pytyp = int
    pindt = pd.Index

    def _asis(self, obj):
        if isinstance(obj.index, pdInt64Index):
            return obj
        self._notimp_error()
    def _drop(self, obj):
        obj = obj.reset_index(drop=True)
        return obj

class StrIndexImp(IndexImplementer):
    """
    Implements a pandas Index consisting of string objects.

    NotImplemented, yet.
    """
    sqlatyp = sqla.String
    pytyp = str
    pindt = pd.Index

    def __init__(self, dfors, case, kwargs):
        raise NotImplementedError()

def _pred(aclass):
    """
    :param aclass
    :return: boolean
    """
    isaclass = inspect.isclass(aclass)
    return isaclass and aclass.__module__ == _pred.__module__

classes = inspect.getmembers(sys.modules[__name__], _pred)

indexingtypes = {cls[0]: cls[1] for cls in classes}
<commit_msg>Fix up constructors on Non-implemented<commit_after>import inspect
import sys

import pandas as pd

pdDatetimeIndex = pd.tseries.index.DatetimeIndex
pdInt64Index = pd.core.index.Int64Index
pdCoreIndex = pd.core.index.Index

import sqlalchemy as sqla

import datetime as dt

class IndexImplementer(object):

    """
    IndexImplementer is the base required to implement
    an index of a specific type.  The
    same instance is created at two points in
    the Trump dataflow:

    1. the datatable getting cached and
    2. the data being served.

    The IndexImplementer should be
    indempotent, and dataframe/series agnostic.
    """

    sqlatyp = sqla.Integer
    pytyp = int
    pindt = pd.Index

    def __init__(self, case, **kwargs):
        """
        :param case: str
            This should match a case used to switch
            the logic created in each subclass of IndexImplementer
        :param kwargs: dict
        """

        self.case = case
        self.k = kwargs
    def orfs_ind_from_str(self, userinput):
        ui = {}
        exec("ui = " + userinput)
        obj = self.pytyp(**ui)
        return obj
    def create_empty(self):
        return self.pindt([])
    def build_ordf(self, orind, orval, colname):
        ordf = pd.DataFrame(index=orind, data=orval, columns=[colname])
        return self._common_passthrough(ordf)
    def _common_passthrough(self, obj):
        try:
            return getattr(self, '_' + self.case)(obj)
        except AttributeError:
            self._notimp_error()
    def _notimp_error(self):
        msg = "Indexing case '{}' unsupported for Index Implementer {}"
        msg = msg.format(self.case, self.__class__.__name__)
        raise NotImplementedError(msg)
    def _notimpobj_error(self, obj):
        msg = "Object of type '{}' unsupported for Index Implementer {}, case {}"
        msg = msg.format(type(obj), self.__class__.__name__, self.case)
        raise TypeError(msg)

    def process_post_db(self, df):
        return self._common_passthrough(df)
    def process_post_feed_cache(self, s):
        return self._common_passthrough(s)
    def process_post_orfs(self, df):
        return self._common_passthrough(df)

class DateIndexImp(IndexImplementer):
    sqlatyp = sqla.Date
    pytyp = dt.date
    pindt = pd.Index
    def __init__(self, case, **kwargs):
        raise NotImplemented()

class DatetimeIndexImp(IndexImplementer):
    """
    Implements a pandas DatetimeIndex

    Cases include:

    * **asis** - Cache timestamps to the database and drop
      any intelligence associated with the index, such as frequency.
      serve a Series with a DatetimeIndex, without frequency.

      If the index consists of 4-digit integers, it will
      be treated as the year, in a date which is of the form
      YYYY-12-31.

    * **asfreq** - Apply 'asfreq' logic prior to cache, and
      apply the same logic when serving.

    * **date_range** - Create a new index, using pandas date_range(), at
      time of cache...  NotImplemented yet.

    * **guess** - NotImplemented yet.
      Attempt to guess the frequency at time of cache,
      and time of serve.

    * **guess_post** - NotImplemented yet.
      Attempt to guess the frequency at time of serve,
      but store the cache unsaved.

    In the event that case hasn't implemented the logic
    to handle a specific datatype, a rudimentary
    attempt to convert it to a DatetimeIndex is applied
    by inspecting the start and end, with the kwargs. passed
    pandas.DatetimeIndex constructor.

    """
    sqlatyp = sqla.DateTime
    pytyp = dt.datetime
    pindt = pd.DatetimeIndex

    def _asis(self, obj):
        if isinstance(obj.index, pdDatetimeIndex):
            return obj
        elif isinstance(obj.index, pdInt64Index):
            if all(len(str(i)) == 4 for i in obj.index):
                #safe to assume we meant years...
                newind = [dt.date(y, 12, 31) for y in obj.index]
                obj.index = newind
                obj.index = obj.index.to_datetime()
                return obj
            self._notimpobj_error(obj)
        elif isinstance(obj.index, pdCoreIndex):
            if isinstance(obj.index[0], dt.date):
                newind = pdDatetimeIndex(obj.index)
                obj.index = newind
                return obj
            self._notimpobj_error(obj)
        else:
            return self._default(obj, **self.k)

    def _asfreq(self, obj):
        if isinstance(obj.index, pdDatetimeIndex):
            obj = obj.asfreq(**self.k)
            return obj
        elif isinstance(obj.index, pdInt64Index):
            if all(len(str(i)) == 4 for i in obj.index):
                #safe to assume we meant years...
                newind = [dt.date(y, 12, 31) for y in obj.index]
                obj.index = newind
                obj.index = obj.index.to_datetime()
                obj = obj.asfreq(**self.k)
                return obj
            self._notimpobj_error(obj)
        elif isinstance(obj.index[0], (str, unicode)):
            newind = pd.DatetimeIndex(obj.index)
            obj.index = newind
            obj = obj.asfreq(**self.k)
            return obj
        else:
            return self._default(obj, **self.k)

    @staticmethod
    def _default(obj, **kwargs):
        start = pd.to_datetime(obj.index[0])
        end = pd.to_datetime(obj.index[-1])
        newind = pdDatetimeIndex(start=start, end=end, **kwargs)
        obj = obj.reindex(newind)
        return obj

class PeriodIndexImp(IndexImplementer):
    """
    Implements a pandas PeriodIndex

    NotImplemented, yet.
    """
    sqlatyp = sqla.DateTime
    pytyp = pd.Period
    pindt = pd.PeriodIndex
    def __init__(self, case, **kwargs):
        raise NotImplementedError()

class IntIndexImp(IndexImplementer):
    """
    Implements a pandas Int64Index.

    Cases include:

    * **asis** - attempts to pass the index through,
      without applying any logic.  Use this, if the index
      is already integers, or unique and integer-like.

    * **drop** - will drop the pandas index, to reset it.

    """

    sqlatyp = sqla.Integer
    pytyp = int
    pindt = pd.Index

    def _asis(self, obj):
        if isinstance(obj.index, pdInt64Index):
            return obj
        self._notimp_error()
    def _drop(self, obj):
        obj = obj.reset_index(drop=True)
        return obj

class StrIndexImp(IndexImplementer):
    """
    Implements a pandas Index consisting of string objects.

    Only method, is "asis"
    """
    sqlatyp = sqla.String
    pytyp = str
    pindt = pd.Index

    def _asis(self, obj):
        if isinstance(obj.index, pdCoreIndex):
            return obj
        self._notimp_error()

def _pred(aclass):
    """
    :param aclass
    :return: boolean
    """
    isaclass = inspect.isclass(aclass)
    return isaclass and aclass.__module__ == _pred.__module__

classes = inspect.getmembers(sys.modules[__name__], _pred)

indexingtypes = {cls[0]: cls[1] for cls in classes}
<|endoftext|>


<commit_before><commit_msg>this isn't really a python module<commit_after><|endoftext|>


<commit_before>import logging

from tsstats.exceptions import InvalidLog
from collections import MutableMapping

logger = logging.getLogger('tsstats')

class Clients(MutableMapping):
    '''
    A high-level-interface to multiple Client-objects
    '''
    def __init__(self, ident_map=None, *args, **kwargs):
        '''
        Initialize a new Client-collection

        :param ident_map: Identity-map (see :ref:`IdentMap`)
        :type ident_map: dict
        '''
        self.ident_map = ident_map or {}

        self.store = dict()
        self.update(dict(*args, **kwargs))

    def __add__(self, client):
        '''
        Add a Client to the collection

        :param client: Client to add to the collection
        :type id_or_uid: Client
        '''
        self.store[client.identifier] = client
        return self

    def __iter__(self):
        '''
        Yield all Client-objects from the collection
        '''
        return iter(self.store.values())

    def __getitem__(self, key):
        return self.store[self.ident_map.get(key, key)]

    def __delitem__(self, key):
        del self.store[key]

    def __len__(self):
        return len(self.store)

    def __setitem__(self, key, value):
        self.store[key] = value

class Client(object):
    '''
    Client provides high-level-access to a Teamspeak-Client
    '''

    def __init__(self, identifier):
        '''
        Initialize a new Client

        :param identifier: Identifier of the client
        :type identifier: int or str
        '''
        # public
        self.identifier = identifier
        self.nick = None
        self.connected = 0
        self.onlinetime = 0
        self.kicks = 0
        self.pkicks = 0
        self.bans = 0
        self.pbans = 0
        self.last_seen = 0
        # private
        self._last_connect = 0

    def connect(self, timestamp):
        '''
        Connect client at `timestamp`

        :param timestamp: time of connect
        :type timestamp: int
        '''
        logger.debug('CONNECT %s', self)
        self.connected += 1
        self._last_connect = timestamp

    def disconnect(self, timestamp):
        '''
        Disconnect client at `timestamp`

        :param timestamp: time of disconnect
        :type timestamp: int
        '''
        logger.debug('DISCONNECT %s', self)
        if not self.connected:
            logger.debug('^ disconnect before connect')
            raise InvalidLog('disconnect before connect!')
        self.connected -= 1
        session_time = timestamp - self._last_connect
        self.onlinetime += session_time
        self.last_seen = timestamp

    def kick(self, target):
        '''
        Let client kick `target`

        :param target: client to kick
        :type target: Client
        '''
        logger.debug('KICK %s -> %s', self, target)
        target.pkicks += 1
        self.kicks += 1

    def ban(self, target):
        '''
        Let client ban `target`

        :param target: client to ban
        :type target: Client
        '''
        logger.debug('BAN %s -> %s', self, target)
        target.pbans += 1
        self.bans += 1

    def __str__(self):
        return '<{},{}>'.format(self.identifier, self.nick)

    def __getitem__(self, item):
        '''
        DEPRECATED
        '''
        return {
            'identifier': self.identifier,
            'nick': self.nick,
            'connected': self.connected,
            'onlinetime': self.onlinetime,
            'kicks': self.kicks,
            'pkicks': self.pkicks,
            'bans': self.bans,
            'pbans': self.pbans,
        }[item]
<commit_msg>add optional nick-arg to tsstats.client.Client.__init__<commit_after>import logging

from tsstats.exceptions import InvalidLog
from collections import MutableMapping

logger = logging.getLogger('tsstats')

class Clients(MutableMapping):
    '''
    A high-level-interface to multiple Client-objects
    '''
    def __init__(self, ident_map=None, *args, **kwargs):
        '''
        Initialize a new Client-collection

        :param ident_map: Identity-map (see :ref:`IdentMap`)
        :type ident_map: dict
        '''
        self.ident_map = ident_map or {}

        self.store = dict()
        self.update(dict(*args, **kwargs))

    def __add__(self, client):
        '''
        Add a Client to the collection

        :param client: Client to add to the collection
        :type id_or_uid: Client
        '''
        self.store[client.identifier] = client
        return self

    def __iter__(self):
        '''
        Yield all Client-objects from the collection
        '''
        return iter(self.store.values())

    def __getitem__(self, key):
        return self.store[self.ident_map.get(key, key)]

    def __delitem__(self, key):
        del self.store[key]

    def __len__(self):
        return len(self.store)

    def __setitem__(self, key, value):
        self.store[key] = value

class Client(object):
    '''
    Client provides high-level-access to a Teamspeak-Client
    '''

    def __init__(self, identifier, nick=None):
        '''
        Initialize a new Client

        :param identifier: Identifier of the client
        :type identifier: int or str
        '''
        # public
        self.identifier = identifier
        self.nick = nick
        self.connected = 0
        self.onlinetime = 0
        self.kicks = 0
        self.pkicks = 0
        self.bans = 0
        self.pbans = 0
        self.last_seen = 0
        # private
        self._last_connect = 0

    def connect(self, timestamp):
        '''
        Connect client at `timestamp`

        :param timestamp: time of connect
        :type timestamp: int
        '''
        logger.debug('CONNECT %s', self)
        self.connected += 1
        self._last_connect = timestamp

    def disconnect(self, timestamp):
        '''
        Disconnect client at `timestamp`

        :param timestamp: time of disconnect
        :type timestamp: int
        '''
        logger.debug('DISCONNECT %s', self)
        if not self.connected:
            logger.debug('^ disconnect before connect')
            raise InvalidLog('disconnect before connect!')
        self.connected -= 1
        session_time = timestamp - self._last_connect
        self.onlinetime += session_time
        self.last_seen = timestamp

    def kick(self, target):
        '''
        Let client kick `target`

        :param target: client to kick
        :type target: Client
        '''
        logger.debug('KICK %s -> %s', self, target)
        target.pkicks += 1
        self.kicks += 1

    def ban(self, target):
        '''
        Let client ban `target`

        :param target: client to ban
        :type target: Client
        '''
        logger.debug('BAN %s -> %s', self, target)
        target.pbans += 1
        self.bans += 1

    def __str__(self):
        return '<{},{}>'.format(self.identifier, self.nick)

    def __getitem__(self, item):
        '''
        DEPRECATED
        '''
        return {
            'identifier': self.identifier,
            'nick': self.nick,
            'connected': self.connected,
            'onlinetime': self.onlinetime,
            'kicks': self.kicks,
            'pkicks': self.pkicks,
            'bans': self.bans,
            'pbans': self.pbans,
        }[item]
<|endoftext|>


<commit_before># ASN.1 types tags
try:
    from sys import version_info
except ImportError:
    version_info = (0, 0)   # a really early version
from operator import getslice
from types import SliceType
from string import join
from pyasn1 import error

tagClassUniversal = 0x00
tagClassApplication = 0x40
tagClassContext = 0x80
tagClassPrivate = 0xC0

tagFormatSimple = 0x00
tagFormatConstructed = 0x20

tagCategoryImplicit = 0x01
tagCategoryExplicit = 0x02
tagCategoryUntagged = 0x04

class Tag:
    def __init__(self, tagClass, tagFormat, tagId):
        self.__tag = (tagClass, tagFormat, tagId)
        self.__uniqTag = (tagClass, tagId)
        self.__hashedUniqTag = hash(self.__uniqTag)
    def __repr__(self):
        return '%s(tagClass=%s, tagFormat=%s, tagId=%s)' % (
            (self.__class__.__name__,) + self.__tag
            )
    def __cmp__(self, other): return cmp(self.__uniqTag, other)
    def __eq__(self, other):
        return self is other or self.__hashedUniqTag == other
    def __hash__(self): return self.__hashedUniqTag
    def __getitem__(self, idx): return self.__tag[idx]
    def __and__(self, (tagClass, tagFormat, tagId)):
        return self.__class__(
            self.__tag&tagClass, self.__tag&tagFormat, self.__tag&tagId
            )
    def __or__(self, (tagClass, tagFormat, tagId)):
        return self.__class__(
            self.__tag[0]|tagClass, self.__tag[1]|tagFormat, self.__tag[2]|tagId
            )

class TagSet:
    def __init__(self, baseTag=(), *superTags):
        self.__baseTag = baseTag
        self.__superTags = superTags
        self.__hashedSuperTags = hash(superTags)
    def __repr__(self):
        return '%s(%s)' % (
            self.__class__.__name__,
            join(map(lambda x: repr(x), self.__superTags), ', ')
            )

    def __add__(self, superTag):
        return apply(
            self.__class__, (self.__baseTag,) + self.__superTags + (superTag,)
            )
    def __radd__(self, superTag):
        return apply(
            self.__class__, (self.__baseTag, superTag) + self.__superTags
            )

    def tagExplicitly(self, superTag):
        tagClass, tagFormat, tagId = superTag
        if tagClass == tagClassUniversal:
            raise error.BadArgumentError(
                'Can\'t tag with UNIVERSAL-class tag'
                )
        if tagFormat != tagFormatConstructed:
            superTag = Tag(tagClass, tagFormatConstructed, tagId)
        return self + superTag

    def tagImplicitly(self, superTag):
        tagClass, tagFormat, tagId = superTag
        if self.__superTags:
            superTag = Tag(tagClass, self.__superTags[-1][1], tagId)
        return self[:-1] + superTag

    def getBaseTag(self): return self.__baseTag
    def __getitem__(self, idx):
        if type(idx) == SliceType:
            return apply(self.__class__,
                         (self.__baseTag,) + getslice(self.__superTags,
                                                      idx.start, idx.stop))
        return self.__superTags[idx]
    if version_info < (2, 0):
        def __getslice__(self, i, j):
            return self[max(0, i):max(0, j):]
    def __cmp__(self, other): return cmp(self.__superTags, other)
    def __eq__(self, other):
        return self is other or self.__hashedSuperTags == other
    def __hash__(self): return self.__hashedSuperTags
    def __len__(self): return len(self.__superTags)
    def isSuperTagSetOf(self, tagSet):
        # XXX optimization?
        l = len(self.__superTags)
        if len(tagSet) < l:
            return
        idx = 0
        while idx < l:
            if self.__superTags[idx] != tagSet[idx]:
                return
            idx = idx + 1
        return 1

def initTagSet(tag): return TagSet(tag, tag)
<commit_msg>minor performance optimization of TagSet.isSuperTagSetOf() method<commit_after># ASN.1 types tags
try:
    from sys import version_info
except ImportError:
    version_info = (0, 0)   # a really early version
from operator import getslice
from types import SliceType
from string import join
from pyasn1 import error

tagClassUniversal = 0x00
tagClassApplication = 0x40
tagClassContext = 0x80
tagClassPrivate = 0xC0

tagFormatSimple = 0x00
tagFormatConstructed = 0x20

tagCategoryImplicit = 0x01
tagCategoryExplicit = 0x02
tagCategoryUntagged = 0x04

class Tag:
    def __init__(self, tagClass, tagFormat, tagId):
        self.__tag = (tagClass, tagFormat, tagId)
        self.__uniqTag = (tagClass, tagId)
        self.__hashedUniqTag = hash(self.__uniqTag)
    def __repr__(self):
        return '%s(tagClass=%s, tagFormat=%s, tagId=%s)' % (
            (self.__class__.__name__,) + self.__tag
            )
    def __cmp__(self, other): return cmp(self.__uniqTag, other)
    def __eq__(self, other):
        return self is other or self.__hashedUniqTag == other
    def __hash__(self): return self.__hashedUniqTag
    def __getitem__(self, idx): return self.__tag[idx]
    def __and__(self, (tagClass, tagFormat, tagId)):
        return self.__class__(
            self.__tag&tagClass, self.__tag&tagFormat, self.__tag&tagId
            )
    def __or__(self, (tagClass, tagFormat, tagId)):
        return self.__class__(
            self.__tag[0]|tagClass, self.__tag[1]|tagFormat, self.__tag[2]|tagId
            )

class TagSet:
    def __init__(self, baseTag=(), *superTags):
        self.__baseTag = baseTag
        self.__superTags = superTags
        self.__hashedSuperTags = hash(superTags)
        self.__lenOfSuperTags = len(superTags)

    def __repr__(self):
        return '%s(%s)' % (
            self.__class__.__name__,
            join(map(lambda x: repr(x), self.__superTags), ', ')
            )

    def __add__(self, superTag):
        return apply(
            self.__class__, (self.__baseTag,) + self.__superTags + (superTag,)
            )
    def __radd__(self, superTag):
        return apply(
            self.__class__, (self.__baseTag, superTag) + self.__superTags
            )

    def tagExplicitly(self, superTag):
        tagClass, tagFormat, tagId = superTag
        if tagClass == tagClassUniversal:
            raise error.BadArgumentError(
                'Can\'t tag with UNIVERSAL-class tag'
                )
        if tagFormat != tagFormatConstructed:
            superTag = Tag(tagClass, tagFormatConstructed, tagId)
        return self + superTag

    def tagImplicitly(self, superTag):
        tagClass, tagFormat, tagId = superTag
        if self.__superTags:
            superTag = Tag(tagClass, self.__superTags[-1][1], tagId)
        return self[:-1] + superTag

    def getBaseTag(self): return self.__baseTag
    def __getitem__(self, idx):
        if type(idx) == SliceType:
            return apply(self.__class__,
                         (self.__baseTag,) + getslice(self.__superTags,
                                                      idx.start, idx.stop))
        return self.__superTags[idx]
    if version_info < (2, 0):
        def __getslice__(self, i, j):
            return self[max(0, i):max(0, j):]
    def __cmp__(self, other): return cmp(self.__superTags, other)
    def __eq__(self, other):
        return self is other or self.__hashedSuperTags == other
    def __hash__(self): return self.__hashedSuperTags
    def __len__(self): return self.__lenOfSuperTags
    def isSuperTagSetOf(self, tagSet):
        if len(tagSet) < self.__lenOfSuperTags:
            return
        idx = self.__lenOfSuperTags - 1
        while idx >= 0:
            if self.__superTags[idx] != tagSet[idx]:
                return
            idx = idx - 1
        return 1

def initTagSet(tag): return TagSet(tag, tag)
<|endoftext|>


<commit_before>#!/usr/bin/python
# @lint-avoid-python-3-compatibility-imports
#
# biosnoop  Trace block device I/O and print details including issuing PID.
#           For Linux, uses BCC, eBPF.
#
# This uses in-kernel eBPF maps to cache process details (PID and comm) by I/O
# request, as well as a starting timestamp for calculating I/O latency.
#
# Copyright (c) 2015 Brendan Gregg.
# Licensed under the Apache License, Version 2.0 (the "License")
#
# 16-Sep-2015   Brendan Gregg   Created this.
# 11-Feb-2016   Allan McAleavy  updated for BPF_PERF_OUTPUT
# 21-Jun-2022   Rocky Xing      Added disk filter support.

from __future__ import print_function
from bcc import BPF
import argparse
import os

# arguments
examples = """examples:
    ./biosnoop           # trace all block I/O
    ./biosnoop -Q        # include OS queued time
    ./biolatency -d sdc  # trace sdc only
"""
parser = argparse.ArgumentParser(
    description="Trace block I/O",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog=examples)
parser.add_argument("-Q", "--queue", action="store_true",
    help="include OS queued time")
parser.add_argument("-d", "--disk", type=str,
    help="Trace this disk only")
parser.add_argument("--ebpf", action="store_true",
    help=argparse.SUPPRESS)
args = parser.parse_args()
debug = 0

# define BPF program
bpf_text="""
#include <uapi/linux/ptrace.h>
#include <linux/blk-mq.h>

// for saving the timestamp and __data_len of each request
struct start_req_t {
    u64 ts;
    u64 data_len;
};

struct val_t {
    u64 ts;
    u32 pid;
    char name[TASK_COMM_LEN];
};

struct data_t {
    u32 pid;
    u64 rwflag;
    u64 delta;
    u64 qdelta;
    u64 sector;
    u64 len;
    u64 ts;
    char disk_name[DISK_NAME_LEN];
    char name[TASK_COMM_LEN];
};

BPF_HASH(start, struct request *, struct start_req_t);
BPF_HASH(infobyreq, struct request *, struct val_t);
BPF_PERF_OUTPUT(events);

// cache PID and comm by-req
int trace_pid_start(struct pt_regs *ctx, struct request *req)
{
    DISK_FILTER

    struct val_t val = {};
    u64 ts;

    if (bpf_get_current_comm(&val.name, sizeof(val.name)) == 0) {
        val.pid = bpf_get_current_pid_tgid() >> 32;
        if (##QUEUE##) {
            val.ts = bpf_ktime_get_ns();
        }
        infobyreq.update(&req, &val);
    }
    return 0;
}

// time block I/O
int trace_req_start(struct pt_regs *ctx, struct request *req)
{
    DISK_FILTER

    struct start_req_t start_req = {
        .ts = bpf_ktime_get_ns(),
        .data_len = req->__data_len
    };
    start.update(&req, &start_req);
    return 0;
}

// output
int trace_req_completion(struct pt_regs *ctx, struct request *req)
{
    struct start_req_t *startp;
    struct val_t *valp;
    struct data_t data = {};
    u64 ts;

    // fetch timestamp and calculate delta
    startp = start.lookup(&req);
    if (startp == 0) {
        // missed tracing issue
        return 0;
    }
    ts = bpf_ktime_get_ns();
    data.delta = ts - startp->ts;
    data.ts = ts / 1000;
    data.qdelta = 0;

    valp = infobyreq.lookup(&req);
    data.len = startp->data_len;
    if (valp == 0) {
        data.name[0] = '?';
        data.name[1] = 0;
    } else {
        if (##QUEUE##) {
            data.qdelta = startp->ts - valp->ts;
        }
        data.pid = valp->pid;
        data.sector = req->__sector;
        bpf_probe_read_kernel(&data.name, sizeof(data.name), valp->name);
        struct gendisk *rq_disk = req->__RQ_DISK__;
        bpf_probe_read_kernel(&data.disk_name, sizeof(data.disk_name),
                       rq_disk->disk_name);
    }

/*
 * The following deals with a kernel version change (in mainline 4.7, although
 * it may be backported to earlier kernels) with how block request write flags
 * are tested. We handle both pre- and post-change versions here. Please avoid
 * kernel version tests like this as much as possible: they inflate the code,
 * test, and maintenance burden.
 */
#ifdef REQ_WRITE
    data.rwflag = !!(req->cmd_flags & REQ_WRITE);
#elif defined(REQ_OP_SHIFT)
    data.rwflag = !!((req->cmd_flags >> REQ_OP_SHIFT) == REQ_OP_WRITE);
#else
    data.rwflag = !!((req->cmd_flags & REQ_OP_MASK) == REQ_OP_WRITE);
#endif

    events.perf_submit(ctx, &data, sizeof(data));
    start.delete(&req);
    infobyreq.delete(&req);

    return 0;
}
"""
if args.queue:
    bpf_text = bpf_text.replace('##QUEUE##', '1')
else:
    bpf_text = bpf_text.replace('##QUEUE##', '0')
if BPF.kernel_struct_has_field(b'request', b'rq_disk') == 1:
    bpf_text = bpf_text.replace('__RQ_DISK__', 'rq_disk')
else:
    bpf_text = bpf_text.replace('__RQ_DISK__', 'q->disk')

if args.disk is not None:
    disk_path = os.path.join('/dev', args.disk)
    if not os.path.exists(disk_path):
        print("no such disk '%s'" % args.disk)
        exit(1)

    stat_info = os.stat(disk_path)
    major = os.major(stat_info.st_rdev)
    minor = os.minor(stat_info.st_rdev)

    disk_field_str = ""
    if BPF.kernel_struct_has_field(b'request', b'rq_disk') == 1:
        disk_field_str = 'req->rq_disk'
    else:
        disk_field_str = 'req->q->disk'

    disk_filter_str = """
    struct gendisk *disk = %s;
    if (!(disk->major == %d && disk->first_minor == %d)) {
        return 0;
    }
    """ % (disk_field_str, major, minor)

    bpf_text = bpf_text.replace('DISK_FILTER', disk_filter_str)
else:
    bpf_text = bpf_text.replace('DISK_FILTER', '')

if debug or args.ebpf:
    print(bpf_text)
    if args.ebpf:
        exit()

# initialize BPF
b = BPF(text=bpf_text)
if BPF.get_kprobe_functions(b'__blk_account_io_start'):
    b.attach_kprobe(event="__blk_account_io_start", fn_name="trace_pid_start")
else:
    b.attach_kprobe(event="blk_account_io_start", fn_name="trace_pid_start")
if BPF.get_kprobe_functions(b'blk_start_request'):
    b.attach_kprobe(event="blk_start_request", fn_name="trace_req_start")
b.attach_kprobe(event="blk_mq_start_request", fn_name="trace_req_start")
if BPF.get_kprobe_functions(b'__blk_account_io_done'):
    b.attach_kprobe(event="__blk_account_io_done", fn_name="trace_req_completion")
else:
    b.attach_kprobe(event="blk_account_io_done", fn_name="trace_req_completion")

# header
print("%-11s %-14s %-7s %-9s %-1s %-10s %-7s" % ("TIME(s)", "COMM", "PID",
    "DISK", "T", "SECTOR", "BYTES"), end="")
if args.queue:
    print("%7s " % ("QUE(ms)"), end="")
print("%7s" % "LAT(ms)")

rwflg = ""
start_ts = 0
prev_ts = 0
delta = 0

# process event
def print_event(cpu, data, size):
    event = b["events"].event(data)

    global start_ts
    if start_ts == 0:
        start_ts = event.ts

    if event.rwflag == 1:
        rwflg = "W"
    else:
        rwflg = "R"

    delta = float(event.ts) - start_ts

    disk_name = event.disk_name.decode('utf-8', 'replace')
    if not disk_name:
        disk_name = '<unknown>'

    print("%-11.6f %-14.14s %-7s %-9s %-1s %-10s %-7s" % (
        delta / 1000000, event.name.decode('utf-8', 'replace'), event.pid,
        disk_name, rwflg, event.sector, event.len), end="")
    if args.queue:
        print("%7.2f " % (float(event.qdelta) / 1000000), end="")
    print("%7.2f" % (float(event.delta) / 1000000))

# loop with callback to print_event
b["events"].open_perf_buffer(print_event, page_cnt=64)
while 1:
    try:
        b.perf_buffer_poll()
    except KeyboardInterrupt:
        exit()
<commit_msg>Update biosnoop.py (#4124)<commit_after>#!/usr/bin/python
# @lint-avoid-python-3-compatibility-imports
#
# biosnoop  Trace block device I/O and print details including issuing PID.
#           For Linux, uses BCC, eBPF.
#
# This uses in-kernel eBPF maps to cache process details (PID and comm) by I/O
# request, as well as a starting timestamp for calculating I/O latency.
#
# Copyright (c) 2015 Brendan Gregg.
# Licensed under the Apache License, Version 2.0 (the "License")
#
# 16-Sep-2015   Brendan Gregg   Created this.
# 11-Feb-2016   Allan McAleavy  updated for BPF_PERF_OUTPUT
# 21-Jun-2022   Rocky Xing      Added disk filter support.

from __future__ import print_function
from bcc import BPF
import argparse
import os

# arguments
examples = """examples:
    ./biosnoop           # trace all block I/O
    ./biosnoop -Q        # include OS queued time
    ./biolatency -d sdc  # trace sdc only
"""
parser = argparse.ArgumentParser(
    description="Trace block I/O",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog=examples)
parser.add_argument("-Q", "--queue", action="store_true",
    help="include OS queued time")
parser.add_argument("-d", "--disk", type=str,
    help="Trace this disk only")
parser.add_argument("--ebpf", action="store_true",
    help=argparse.SUPPRESS)
args = parser.parse_args()
debug = 0

# define BPF program
bpf_text="""
#include <uapi/linux/ptrace.h>
#include <linux/blk-mq.h>

// for saving the timestamp and __data_len of each request
struct start_req_t {
    u64 ts;
    u64 data_len;
};

struct val_t {
    u64 ts;
    u32 pid;
    char name[TASK_COMM_LEN];
};

struct data_t {
    u32 pid;
    u64 rwflag;
    u64 delta;
    u64 qdelta;
    u64 sector;
    u64 len;
    u64 ts;
    char disk_name[DISK_NAME_LEN];
    char name[TASK_COMM_LEN];
};

BPF_HASH(start, struct request *, struct start_req_t);
BPF_HASH(infobyreq, struct request *, struct val_t);
BPF_PERF_OUTPUT(events);

// cache PID and comm by-req
int trace_pid_start(struct pt_regs *ctx, struct request *req)
{
    DISK_FILTER

    struct val_t val = {};
    u64 ts;

    if (bpf_get_current_comm(&val.name, sizeof(val.name)) == 0) {
        val.pid = bpf_get_current_pid_tgid() >> 32;
        if (##QUEUE##) {
            val.ts = bpf_ktime_get_ns();
        }
        infobyreq.update(&req, &val);
    }
    return 0;
}

// time block I/O
int trace_req_start(struct pt_regs *ctx, struct request *req)
{
    DISK_FILTER

    struct start_req_t start_req = {
        .ts = bpf_ktime_get_ns(),
        .data_len = req->__data_len
    };
    start.update(&req, &start_req);
    return 0;
}

// output
int trace_req_completion(struct pt_regs *ctx, struct request *req)
{
    struct start_req_t *startp;
    struct val_t *valp;
    struct data_t data = {};
    u64 ts;

    // fetch timestamp and calculate delta
    startp = start.lookup(&req);
    if (startp == 0) {
        // missed tracing issue
        return 0;
    }
    ts = bpf_ktime_get_ns();
    data.delta = ts - startp->ts;
    data.ts = ts / 1000;
    data.qdelta = 0;

    valp = infobyreq.lookup(&req);
    data.len = startp->data_len;
    if (valp == 0) {
        data.name[0] = '?';
        data.name[1] = 0;
    } else {
        if (##QUEUE##) {
            data.qdelta = startp->ts - valp->ts;
        }
        data.pid = valp->pid;
        data.sector = req->__sector;
        bpf_probe_read_kernel(&data.name, sizeof(data.name), valp->name);
        struct gendisk *rq_disk = req->__RQ_DISK__;
        bpf_probe_read_kernel(&data.disk_name, sizeof(data.disk_name),
                       rq_disk->disk_name);
    }

/*
 * The following deals with a kernel version change (in mainline 4.7, although
 * it may be backported to earlier kernels) with how block request write flags
 * are tested. We handle both pre- and post-change versions here. Please avoid
 * kernel version tests like this as much as possible: they inflate the code,
 * test, and maintenance burden.
 */
#ifdef REQ_WRITE
    data.rwflag = !!(req->cmd_flags & REQ_WRITE);
#elif defined(REQ_OP_SHIFT)
    data.rwflag = !!((req->cmd_flags >> REQ_OP_SHIFT) == REQ_OP_WRITE);
#else
    data.rwflag = !!((req->cmd_flags & REQ_OP_MASK) == REQ_OP_WRITE);
#endif

    events.perf_submit(ctx, &data, sizeof(data));
    start.delete(&req);
    infobyreq.delete(&req);

    return 0;
}
"""
if args.queue:
    bpf_text = bpf_text.replace('##QUEUE##', '1')
else:
    bpf_text = bpf_text.replace('##QUEUE##', '0')
if BPF.kernel_struct_has_field(b'request', b'rq_disk') == 1:
    bpf_text = bpf_text.replace('__RQ_DISK__', 'rq_disk')
else:
    bpf_text = bpf_text.replace('__RQ_DISK__', 'q->disk')

if args.disk is not None:
    disk_path = os.path.join('/dev', args.disk)
    if not os.path.exists(disk_path):
        print("no such disk '%s'" % args.disk)
        exit(1)

    stat_info = os.stat(disk_path)
    major = os.major(stat_info.st_rdev)
    minor = os.minor(stat_info.st_rdev)

    disk_field_str = ""
    if BPF.kernel_struct_has_field(b'request', b'rq_disk') == 1:
        disk_field_str = 'req->rq_disk'
    else:
        disk_field_str = 'req->q->disk'

    disk_filter_str = """
    struct gendisk *disk = %s;
    if (!(disk->major == %d && disk->first_minor == %d)) {
        return 0;
    }
    """ % (disk_field_str, major, minor)

    bpf_text = bpf_text.replace('DISK_FILTER', disk_filter_str)
else:
    bpf_text = bpf_text.replace('DISK_FILTER', '')

if debug or args.ebpf:
    print(bpf_text)
    if args.ebpf:
        exit()

# initialize BPF
b = BPF(text=bpf_text)
if BPF.get_kprobe_functions(b'__blk_account_io_start'):
    b.attach_kprobe(event="__blk_account_io_start", fn_name="trace_pid_start")
else:
    b.attach_kprobe(event="blk_account_io_start", fn_name="trace_pid_start")
if BPF.get_kprobe_functions(b'blk_start_request'):
    b.attach_kprobe(event="blk_start_request", fn_name="trace_req_start")
else:
    b.attach_kprobe(event="blk_mq_start_request", fn_name="trace_req_start")
if BPF.get_kprobe_functions(b'__blk_account_io_done'):
    b.attach_kprobe(event="__blk_account_io_done", fn_name="trace_req_completion")
else:
    b.attach_kprobe(event="blk_account_io_done", fn_name="trace_req_completion")

# header
print("%-11s %-14s %-7s %-9s %-1s %-10s %-7s" % ("TIME(s)", "COMM", "PID",
    "DISK", "T", "SECTOR", "BYTES"), end="")
if args.queue:
    print("%7s " % ("QUE(ms)"), end="")
print("%7s" % "LAT(ms)")

rwflg = ""
start_ts = 0
prev_ts = 0
delta = 0

# process event
def print_event(cpu, data, size):
    event = b["events"].event(data)

    global start_ts
    if start_ts == 0:
        start_ts = event.ts

    if event.rwflag == 1:
        rwflg = "W"
    else:
        rwflg = "R"

    delta = float(event.ts) - start_ts

    disk_name = event.disk_name.decode('utf-8', 'replace')
    if not disk_name:
        disk_name = '<unknown>'

    print("%-11.6f %-14.14s %-7s %-9s %-1s %-10s %-7s" % (
        delta / 1000000, event.name.decode('utf-8', 'replace'), event.pid,
        disk_name, rwflg, event.sector, event.len), end="")
    if args.queue:
        print("%7.2f " % (float(event.qdelta) / 1000000), end="")
    print("%7.2f" % (float(event.delta) / 1000000))

# loop with callback to print_event
b["events"].open_perf_buffer(print_event, page_cnt=64)
while 1:
    try:
        b.perf_buffer_poll()
    except KeyboardInterrupt:
        exit()
<|endoftext|>