<commit_before># Copyright 2021-2022 Google LLC
#
# 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
#
#      https://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.

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
import asyncio
import logging
import usb1
import threading
import collections
from colors import color

from .common import Transport, ParserSource
from .. import hci

# -----------------------------------------------------------------------------
# Logging
# -----------------------------------------------------------------------------
logger = logging.getLogger(__name__)

# -----------------------------------------------------------------------------
async def open_usb_transport(spec):
    '''
    Open a USB transport.
    The parameter string has this syntax:
    either <index> or
    <vendor>:<product> or
    <vendor>:<product>/<serial-number>] or
    <vendor>:<product>#<index>
    With <index> as the 0-based index to select amongst all the devices that appear
    to be supporting Bluetooth HCI (0 being the first one), or
    Where <vendor> and <product> are the vendor ID and product ID in hexadecimal. The
    /<serial-number> suffix or #<index> suffix max be specified when more than one device with
    the same vendor and product identifiers are present.

    Examples:
    0 --> the first BT USB dongle
    04b4:f901 --> the BT USB dongle with vendor=04b4 and product=f901
    04b4:f901#2 --> the third USB device with vendor=04b4 and product=f901
    04b4:f901/00E04C239987 --> the BT USB dongle with vendor=04b4 and product=f901 and serial number 00E04C239987
    '''

    USB_RECIPIENT_DEVICE                             = 0x00
    USB_REQUEST_TYPE_CLASS                           = 0x01 << 5
    USB_DEVICE_CLASS_WIRELESS_CONTROLLER             = 0xE0
    USB_DEVICE_SUBCLASS_RF_CONTROLLER                = 0x01
    USB_DEVICE_PROTOCOL_BLUETOOTH_PRIMARY_CONTROLLER = 0x01
    USB_ENDPOINT_TRANSFER_TYPE_BULK                  = 0x02
    USB_ENDPOINT_TRANSFER_TYPE_INTERRUPT             = 0x03
    USB_ENDPOINT_IN                                  = 0x80

    READ_SIZE = 1024

    class UsbPacketSink:
        def __init__(self, device, acl_out):
            self.device      = device
            self.acl_out     = acl_out
            self.transfer    = device.getTransfer()
            self.packets     = collections.deque()  # Queue of packets waiting to be sent
            self.loop        = asyncio.get_running_loop()
            self.cancel_done = self.loop.create_future()
            self.closed      = False

        def start(self):
            pass

        def on_packet(self, packet):
            # Ignore packets if we're closed
            if self.closed:
                return

            if len(packet) == 0:
                logger.warning('packet too short')
                return

            # Queue the packet
            self.packets.append(packet)
            if len(self.packets) == 1:
                # The queue was previously empty, re-prime the pump
                self.process_queue()

        def on_packet_sent(self, transfer):
            status = transfer.getStatus()
            # logger.debug(f'<<< USB out transfer callback: status={status}')

            if status == usb1.TRANSFER_COMPLETED:
                self.loop.call_soon_threadsafe(self.on_packet_sent_)
            elif status == usb1.TRANSFER_CANCELLED:
                self.loop.call_soon_threadsafe(self.cancel_done.set_result, None)
            else:
                logger.warning(color(f'!!! out transfer not completed: status={status}', 'red'))

        def on_packet_sent_(self):
            if self.packets:
                self.packets.popleft()
                self.process_queue()

        def process_queue(self):
            if len(self.packets) == 0:
                return  # Nothing to do

            packet = self.packets[0]
            packet_type = packet[0]
            if packet_type == hci.HCI_ACL_DATA_PACKET:
                self.transfer.setBulk(
                    self.acl_out,
                    packet[1:],
                    callback=self.on_packet_sent
                )
                logger.debug('submit ACL')
                self.transfer.submit()
            elif packet_type == hci.HCI_COMMAND_PACKET:
                self.transfer.setControl(
                    USB_RECIPIENT_DEVICE | USB_REQUEST_TYPE_CLASS, 0, 0, 0,
                    packet[1:],
                    callback=self.on_packet_sent
                )
                logger.debug('submit COMMAND')
                self.transfer.submit()
            else:
                logger.warning(color(f'unsupported packet type {packet_type}', 'red'))

        async def close(self):
            self.closed = True

            # Empty the packet queue so that we don't send any more data
            self.packets.clear()

            # If we have a transfer in flight, cancel it
            if self.transfer.isSubmitted():
                # Try to cancel the transfer, but that may fail because it may have already completed
                try:
                    self.transfer.cancel()

                    logger.debug('waiting for OUT transfer cancellation to be done...')
                    await self.cancel_done
                    logger.debug('OUT transfer cancellation done')
                except usb1.USBError:
                    logger.debug('OUT transfer likely already completed')

    class UsbPacketSource(asyncio.Protocol, ParserSource):
        def __init__(self, context, device, acl_in, events_in):
            super().__init__()
            self.context         = context
            self.device          = device
            self.acl_in          = acl_in
            self.events_in       = events_in
            self.loop            = asyncio.get_running_loop()
            self.queue           = asyncio.Queue()
            self.closed          = False
            self.event_loop_done = self.loop.create_future()
            self.cancel_done = {
                hci.HCI_EVENT_PACKET:    self.loop.create_future(),
                hci.HCI_ACL_DATA_PACKET: self.loop.create_future()
            }

            # Create a thread to process events
            self.event_thread = threading.Thread(target=self.run)

        def start(self):
            # Set up transfer objects for input
            self.events_in_transfer = device.getTransfer()
            self.events_in_transfer.setInterrupt(
                self.events_in,
                READ_SIZE,
                callback=self.on_packet_received,
                user_data=hci.HCI_EVENT_PACKET
            )
            self.events_in_transfer.submit()

            self.acl_in_transfer = device.getTransfer()
            self.acl_in_transfer.setBulk(
                self.acl_in,
                READ_SIZE,
                callback=self.on_packet_received,
                user_data=hci.HCI_ACL_DATA_PACKET
            )
            self.acl_in_transfer.submit()

            self.dequeue_task = self.loop.create_task(self.dequeue())
            self.event_thread.start()

        def on_packet_received(self, transfer):
            packet_type = transfer.getUserData()
            status = transfer.getStatus()
            # logger.debug(f'<<< USB IN transfer callback: status={status} packet_type={packet_type} length={transfer.getActualLength()}')

            if status == usb1.TRANSFER_COMPLETED:
                packet = bytes([packet_type]) + transfer.getBuffer()[:transfer.getActualLength()]
                self.loop.call_soon_threadsafe(self.queue.put_nowait, packet)
            elif status == usb1.TRANSFER_CANCELLED:
                self.loop.call_soon_threadsafe(self.cancel_done[packet_type].set_result, None)
                return
            else:
                logger.warning(color(f'!!! transfer not completed: status={status}', 'red'))

            # Re-submit the transfer so we can receive more data
            transfer.submit()

        async def dequeue(self):
            while not self.closed:
                try:
                    packet = await self.queue.get()
                except asyncio.CancelledError:
                    return
                self.parser.feed_data(packet)

        def run(self):
            logger.debug('starting USB event loop')
            while self.events_in_transfer.isSubmitted() or self.acl_in_transfer.isSubmitted():
                try:
                    self.context.handleEvents()
                except usb1.USBErrorInterrupted:
                    pass

            logger.debug('USB event loop done')
            self.loop.call_soon_threadsafe(self.event_loop_done.set_result, None)

        async def close(self):
            self.closed = True
            self.dequeue_task.cancel()

            # Cancel the transfers
            for transfer in (self.events_in_transfer, self.acl_in_transfer):
                if transfer.isSubmitted():
                    # Try to cancel the transfer, but that may fail because it may have already completed
                    packet_type = transfer.getUserData()
                    try:
                        transfer.cancel()
                        logger.debug(f'waiting for IN[{packet_type}] transfer cancellation to be done...')
                        await self.cancel_done[packet_type]
                        logger.debug(f'IN[{packet_type}] transfer cancellation done')
                    except usb1.USBError:
                        logger.debug(f'IN[{packet_type}] transfer likely already completed')

            # Wait for the thread to terminate
            await self.event_loop_done

    class UsbTransport(Transport):
        def __init__(self, context, device, interface, source, sink):
            super().__init__(source, sink)
            self.context   = context
            self.device    = device
            self.interface = interface

            # Get exclusive access
            device.claimInterface(interface)

            # The source and sink can now start
            source.start()
            sink.start()

        async def close(self):
            await self.source.close()
            await self.sink.close()
            self.device.releaseInterface(self.interface)
            self.device.close()
            self.context.close()

    # Find the device according to the spec moniker
    context = usb1.USBContext()
    context.open()
    try:
        found = None
        if ':' in spec:
            vendor_id, product_id = spec.split(':')
            serial_number = None
            device_index = 0
            if '/' in product_id:
                product_id, serial_number = product_id.split('/')
            elif '#' in product_id:
                product_id, device_index_str = product_id.split('#')
                device_index = int(device_index_str)

            for device in context.getDeviceIterator(skip_on_error=True):
                if (
                    device.getVendorID() == int(vendor_id, 16) and
                    device.getProductID() == int(product_id, 16) and
                    (serial_number is None or device.getSerialNumber() == serial_number)
                ):
                    if device_index == 0:
                        found = device
                        break
                    device_index -= 1
                device.close()
        else:
            device_index = int(spec)
            for device in context.getDeviceIterator(skip_on_error=True):
                if (
                    device.getDeviceClass()    == USB_DEVICE_CLASS_WIRELESS_CONTROLLER and
                    device.getDeviceSubClass() == USB_DEVICE_SUBCLASS_RF_CONTROLLER and
                    device.getDeviceProtocol() == USB_DEVICE_PROTOCOL_BLUETOOTH_PRIMARY_CONTROLLER
                ):
                    if device_index == 0:
                        found = device
                        break
                    device_index -= 1
                device.close()

        if found is None:
            context.close()
            raise ValueError('device not found')

        logger.debug(f'USB Device: {found}')

        # Look for an interface and setting with the right types of endpoints
        def find_endpoints(device):
            for (configuration_index, configuration) in enumerate(device):
                interface = None
                for interface in configuration:
                    setting = None
                    for setting in interface:
                        events_in = None
                        acl_in    = None
                        acl_out   = None
                        for endpoint in setting:
                            attributes = endpoint.getAttributes()
                            address    = endpoint.getAddress()
                            if attributes & 0x03 == USB_ENDPOINT_TRANSFER_TYPE_BULK:
                                if address & USB_ENDPOINT_IN and acl_in is None:
                                    acl_in = address
                                elif acl_out is None:
                                    acl_out = address
                            elif attributes & 0x03 == USB_ENDPOINT_TRANSFER_TYPE_INTERRUPT:
                                if address & USB_ENDPOINT_IN and events_in is None:
                                    events_in = address

                        # Return if we found all 3 endpoints
                        if acl_in is not None and acl_out is not None and events_in is not None:
                            return (configuration_index + 1, setting.getNumber(), acl_in, acl_out, events_in)
                        else:
                            logger.debug(f'skipping configuration {configuration_index + 1} / interface {setting.getNumber()}')

        endpoints = find_endpoints(found)
        if endpoints is None:
            logger.warning('no compatible interface found')
            raise ValueError('no compatible interface found for device')
        (configuration, interface, acl_in, acl_out, events_in) = endpoints
        logger.debug(
            f'selected endpoints: configuration={configuration}, '
            f'interface={interface}, '
            f'acl_in=0x{acl_in:02X}, '
            f'acl_out=0x{acl_out:02X}, '
            f'events_in=0x{events_in:02X}, '
        )

        device = found.open()

        # Detach the kernel driver if supported and needed
        if usb1.hasCapability(usb1.CAP_SUPPORTS_DETACH_KERNEL_DRIVER):
            try:
                if device.kernelDriverActive(interface):
                    logger.debug("detaching kernel driver")
                    device.detachKernelDriver(interface)
            except usb1.USBError:
                pass

        # Set the configuration if needed
        try:
            current_configuration = device.getConfiguration()
            logger.debug(f'current configuration = {current_configuration}')
        except usb1.USBError:
            current_configuration = 0

        if current_configuration != configuration:
            try:
                logger.debug(f'setting configuration {configuration}')
                device.setConfiguration(configuration)
            except usb1.USBError:
                logger.warning('failed to set configuration')

        source = UsbPacketSource(context, device, acl_in, events_in)
        sink   = UsbPacketSink(device, acl_out)
        return UsbTransport(context, device, interface, source, sink)
    except usb1.USBError as error:
        logger.warning(color(f'!!! failed to open USB device: {error}', 'red'))
        context.close()
        raise
<commit_msg>add support for alternate settings<commit_after># Copyright 2021-2022 Google LLC
#
# 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
#
#      https://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.

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
import asyncio
import logging
import usb1
import threading
import collections
from colors import color

from .common import Transport, ParserSource
from .. import hci

# -----------------------------------------------------------------------------
# Logging
# -----------------------------------------------------------------------------
logger = logging.getLogger(__name__)

# -----------------------------------------------------------------------------
async def open_usb_transport(spec):
    '''
    Open a USB transport.
    The parameter string has this syntax:
    either <index> or
    <vendor>:<product> or
    <vendor>:<product>/<serial-number>] or
    <vendor>:<product>#<index>
    With <index> as the 0-based index to select amongst all the devices that appear
    to be supporting Bluetooth HCI (0 being the first one), or
    Where <vendor> and <product> are the vendor ID and product ID in hexadecimal. The
    /<serial-number> suffix or #<index> suffix max be specified when more than one device with
    the same vendor and product identifiers are present.

    Examples:
    0 --> the first BT USB dongle
    04b4:f901 --> the BT USB dongle with vendor=04b4 and product=f901
    04b4:f901#2 --> the third USB device with vendor=04b4 and product=f901
    04b4:f901/00E04C239987 --> the BT USB dongle with vendor=04b4 and product=f901 and serial number 00E04C239987
    '''

    USB_RECIPIENT_DEVICE                             = 0x00
    USB_REQUEST_TYPE_CLASS                           = 0x01 << 5
    USB_DEVICE_CLASS_WIRELESS_CONTROLLER             = 0xE0
    USB_DEVICE_SUBCLASS_RF_CONTROLLER                = 0x01
    USB_DEVICE_PROTOCOL_BLUETOOTH_PRIMARY_CONTROLLER = 0x01
    USB_ENDPOINT_TRANSFER_TYPE_BULK                  = 0x02
    USB_ENDPOINT_TRANSFER_TYPE_INTERRUPT             = 0x03
    USB_ENDPOINT_IN                                  = 0x80

    READ_SIZE = 1024

    class UsbPacketSink:
        def __init__(self, device, acl_out):
            self.device      = device
            self.acl_out     = acl_out
            self.transfer    = device.getTransfer()
            self.packets     = collections.deque()  # Queue of packets waiting to be sent
            self.loop        = asyncio.get_running_loop()
            self.cancel_done = self.loop.create_future()
            self.closed      = False

        def start(self):
            pass

        def on_packet(self, packet):
            # Ignore packets if we're closed
            if self.closed:
                return

            if len(packet) == 0:
                logger.warning('packet too short')
                return

            # Queue the packet
            self.packets.append(packet)
            if len(self.packets) == 1:
                # The queue was previously empty, re-prime the pump
                self.process_queue()

        def on_packet_sent(self, transfer):
            status = transfer.getStatus()
            # logger.debug(f'<<< USB out transfer callback: status={status}')

            if status == usb1.TRANSFER_COMPLETED:
                self.loop.call_soon_threadsafe(self.on_packet_sent_)
            elif status == usb1.TRANSFER_CANCELLED:
                self.loop.call_soon_threadsafe(self.cancel_done.set_result, None)
            else:
                logger.warning(color(f'!!! out transfer not completed: status={status}', 'red'))

        def on_packet_sent_(self):
            if self.packets:
                self.packets.popleft()
                self.process_queue()

        def process_queue(self):
            if len(self.packets) == 0:
                return  # Nothing to do

            packet = self.packets[0]
            packet_type = packet[0]
            if packet_type == hci.HCI_ACL_DATA_PACKET:
                self.transfer.setBulk(
                    self.acl_out,
                    packet[1:],
                    callback=self.on_packet_sent
                )
                logger.debug('submit ACL')
                self.transfer.submit()
            elif packet_type == hci.HCI_COMMAND_PACKET:
                self.transfer.setControl(
                    USB_RECIPIENT_DEVICE | USB_REQUEST_TYPE_CLASS, 0, 0, 0,
                    packet[1:],
                    callback=self.on_packet_sent
                )
                logger.debug('submit COMMAND')
                self.transfer.submit()
            else:
                logger.warning(color(f'unsupported packet type {packet_type}', 'red'))

        async def close(self):
            self.closed = True

            # Empty the packet queue so that we don't send any more data
            self.packets.clear()

            # If we have a transfer in flight, cancel it
            if self.transfer.isSubmitted():
                # Try to cancel the transfer, but that may fail because it may have already completed
                try:
                    self.transfer.cancel()

                    logger.debug('waiting for OUT transfer cancellation to be done...')
                    await self.cancel_done
                    logger.debug('OUT transfer cancellation done')
                except usb1.USBError:
                    logger.debug('OUT transfer likely already completed')

    class UsbPacketSource(asyncio.Protocol, ParserSource):
        def __init__(self, context, device, acl_in, events_in):
            super().__init__()
            self.context         = context
            self.device          = device
            self.acl_in          = acl_in
            self.events_in       = events_in
            self.loop            = asyncio.get_running_loop()
            self.queue           = asyncio.Queue()
            self.closed          = False
            self.event_loop_done = self.loop.create_future()
            self.cancel_done = {
                hci.HCI_EVENT_PACKET:    self.loop.create_future(),
                hci.HCI_ACL_DATA_PACKET: self.loop.create_future()
            }

            # Create a thread to process events
            self.event_thread = threading.Thread(target=self.run)

        def start(self):
            # Set up transfer objects for input
            self.events_in_transfer = device.getTransfer()
            self.events_in_transfer.setInterrupt(
                self.events_in,
                READ_SIZE,
                callback=self.on_packet_received,
                user_data=hci.HCI_EVENT_PACKET
            )
            self.events_in_transfer.submit()

            self.acl_in_transfer = device.getTransfer()
            self.acl_in_transfer.setBulk(
                self.acl_in,
                READ_SIZE,
                callback=self.on_packet_received,
                user_data=hci.HCI_ACL_DATA_PACKET
            )
            self.acl_in_transfer.submit()

            self.dequeue_task = self.loop.create_task(self.dequeue())
            self.event_thread.start()

        def on_packet_received(self, transfer):
            packet_type = transfer.getUserData()
            status = transfer.getStatus()
            # logger.debug(f'<<< USB IN transfer callback: status={status} packet_type={packet_type} length={transfer.getActualLength()}')

            if status == usb1.TRANSFER_COMPLETED:
                packet = bytes([packet_type]) + transfer.getBuffer()[:transfer.getActualLength()]
                self.loop.call_soon_threadsafe(self.queue.put_nowait, packet)
            elif status == usb1.TRANSFER_CANCELLED:
                self.loop.call_soon_threadsafe(self.cancel_done[packet_type].set_result, None)
                return
            else:
                logger.warning(color(f'!!! transfer not completed: status={status}', 'red'))

            # Re-submit the transfer so we can receive more data
            transfer.submit()

        async def dequeue(self):
            while not self.closed:
                try:
                    packet = await self.queue.get()
                except asyncio.CancelledError:
                    return
                self.parser.feed_data(packet)

        def run(self):
            logger.debug('starting USB event loop')
            while self.events_in_transfer.isSubmitted() or self.acl_in_transfer.isSubmitted():
                try:
                    self.context.handleEvents()
                except usb1.USBErrorInterrupted:
                    pass

            logger.debug('USB event loop done')
            self.loop.call_soon_threadsafe(self.event_loop_done.set_result, None)

        async def close(self):
            self.closed = True
            self.dequeue_task.cancel()

            # Cancel the transfers
            for transfer in (self.events_in_transfer, self.acl_in_transfer):
                if transfer.isSubmitted():
                    # Try to cancel the transfer, but that may fail because it may have already completed
                    packet_type = transfer.getUserData()
                    try:
                        transfer.cancel()
                        logger.debug(f'waiting for IN[{packet_type}] transfer cancellation to be done...')
                        await self.cancel_done[packet_type]
                        logger.debug(f'IN[{packet_type}] transfer cancellation done')
                    except usb1.USBError:
                        logger.debug(f'IN[{packet_type}] transfer likely already completed')

            # Wait for the thread to terminate
            await self.event_loop_done

    class UsbTransport(Transport):
        def __init__(self, context, device, interface, setting, source, sink):
            super().__init__(source, sink)
            self.context   = context
            self.device    = device
            self.interface = interface

            # Get exclusive access
            device.claimInterface(interface)

            # Set the alternate setting if not the default
            if setting != 0:
                device.setInterfaceAltSetting(interface, setting)

            # The source and sink can now start
            source.start()
            sink.start()

        async def close(self):
            await self.source.close()
            await self.sink.close()
            self.device.releaseInterface(self.interface)
            self.device.close()
            self.context.close()

    # Find the device according to the spec moniker
    context = usb1.USBContext()
    context.open()
    try:
        found = None
        if ':' in spec:
            vendor_id, product_id = spec.split(':')
            serial_number = None
            device_index = 0
            if '/' in product_id:
                product_id, serial_number = product_id.split('/')
            elif '#' in product_id:
                product_id, device_index_str = product_id.split('#')
                device_index = int(device_index_str)

            for device in context.getDeviceIterator(skip_on_error=True):
                if (
                    device.getVendorID() == int(vendor_id, 16) and
                    device.getProductID() == int(product_id, 16) and
                    (serial_number is None or device.getSerialNumber() == serial_number)
                ):
                    if device_index == 0:
                        found = device
                        break
                    device_index -= 1
                device.close()
        else:
            device_index = int(spec)
            for device in context.getDeviceIterator(skip_on_error=True):
                if (
                    device.getDeviceClass()    == USB_DEVICE_CLASS_WIRELESS_CONTROLLER and
                    device.getDeviceSubClass() == USB_DEVICE_SUBCLASS_RF_CONTROLLER and
                    device.getDeviceProtocol() == USB_DEVICE_PROTOCOL_BLUETOOTH_PRIMARY_CONTROLLER
                ):
                    if device_index == 0:
                        found = device
                        break
                    device_index -= 1
                device.close()

        if found is None:
            context.close()
            raise ValueError('device not found')

        logger.debug(f'USB Device: {found}')

        # Look for the first interface with the right class and endpoints
        def find_endpoints(device):
            for (configuration_index, configuration) in enumerate(device):
                interface = None
                for interface in configuration:
                    setting = None
                    for setting in interface:
                        if (
                            setting.getClass() != USB_DEVICE_CLASS_WIRELESS_CONTROLLER or
                            setting.getSubClass() != USB_DEVICE_SUBCLASS_RF_CONTROLLER or
                            setting.getProtocol() != USB_DEVICE_PROTOCOL_BLUETOOTH_PRIMARY_CONTROLLER
                        ):
                            continue

                        events_in = None
                        acl_in    = None
                        acl_out   = None
                        for endpoint in setting:
                            attributes = endpoint.getAttributes()
                            address    = endpoint.getAddress()
                            if attributes & 0x03 == USB_ENDPOINT_TRANSFER_TYPE_BULK:
                                if address & USB_ENDPOINT_IN and acl_in is None:
                                    acl_in = address
                                elif acl_out is None:
                                    acl_out = address
                            elif attributes & 0x03 == USB_ENDPOINT_TRANSFER_TYPE_INTERRUPT:
                                if address & USB_ENDPOINT_IN and events_in is None:
                                    events_in = address

                        # Return if we found all 3 endpoints
                        if acl_in is not None and acl_out is not None and events_in is not None:
                            return (
                                configuration_index + 1,
                                setting.getNumber(),
                                setting.getAlternateSetting(),
                                acl_in,
                                acl_out,
                                events_in
                            )
                        else:
                            logger.debug(f'skipping configuration {configuration_index + 1} / interface {setting.getNumber()}')

        endpoints = find_endpoints(found)
        if endpoints is None:
            raise ValueError('no compatible interface found for device')
        (configuration, interface, setting, acl_in, acl_out, events_in) = endpoints
        logger.debug(
            f'selected endpoints: configuration={configuration}, '
            f'interface={interface}, '
            f'setting={setting}, '
            f'acl_in=0x{acl_in:02X}, '
            f'acl_out=0x{acl_out:02X}, '
            f'events_in=0x{events_in:02X}, '
        )

        device = found.open()

        # Detach the kernel driver if supported and needed
        if usb1.hasCapability(usb1.CAP_SUPPORTS_DETACH_KERNEL_DRIVER):
            try:
                if device.kernelDriverActive(interface):
                    logger.debug("detaching kernel driver")
                    device.detachKernelDriver(interface)
            except usb1.USBError:
                pass

        # Set the configuration if needed
        try:
            current_configuration = device.getConfiguration()
            logger.debug(f'current configuration = {current_configuration}')
        except usb1.USBError:
            current_configuration = 0

        if current_configuration != configuration:
            try:
                logger.debug(f'setting configuration {configuration}')
                device.setConfiguration(configuration)
            except usb1.USBError:
                logger.warning('failed to set configuration')

        source = UsbPacketSource(context, device, acl_in, events_in)
        sink   = UsbPacketSink(device, acl_out)
        return UsbTransport(context, device, interface, setting, source, sink)
    except usb1.USBError as error:
        logger.warning(color(f'!!! failed to open USB device: {error}', 'red'))
        context.close()
        raise
<|endoftext|>


<commit_before>from __future__ import division

import os
import sys
import math
import numpy as np
import cv2
from multiprocessing.pool import ThreadPool
from collections import deque

lk_params = dict(
	winSize = (15,15),
	maxLevel = 2,
	criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03)
)

feature_params = dict(
	maxCorners = 500,
	qualityLevel = 0.3,
	minDistance = 7,
	blockSize = 7
)

def meanshift(p, radius, points):
	while True:
		pold = p
		inliers = (np.linalg.norm(points - p, axis=1) <= radius)
		p = np.mean(points[inliers], axis=0)
		shift = np.linalg.norm(p - pold)
		if shift < radius * 0.1:
			break

	return p

def clustering(radius, points):
	centroids = np.zeros((0,2), dtype=np.float32)
	labels = np.zeros((len(points),), dtype=np.int0)
	labels[:] = -1
	for i,pt in enumerate(points):
		pt1 = meanshift(pt, radius, points)
		# find nearest centroid
		if len(centroids) == 0:
			labels[i] = 0
			centroids = np.vstack([centroids, pt1])
		else:
			d = np.linalg.norm(centroids - pt1, axis=1)
			imin = np.argmin(d)
			if d[imin] < radius:
				labels[i] = imin
			else:
				labels[i] = len(centroids)
				centroids = np.vstack([centroids, pt1])

	return (centroids, labels)

class DummyTask:
    def __init__(self, fn, args):
        self.data = fn(*args)
    def ready(self):
        return True
    def get(self):
        return self.data

def to_f8(array):
	return tuple(np.rint(array * 256).astype('int'))

def scale_frame(frame):
	frame = cv2.resize(frame, None, fx=framescale, fy=framescale, interpolation=cv2.INTER_AREA)
	framegray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
	return framegray

def flow_frame(tracks, prevframe, curframe):
	assert prevframe is not None and curframe is not None

	vis = cv2.cvtColor(curframe, cv2.COLOR_GRAY2BGR)

	if len(tracks) > 0:
		if 0:
			p0 = np.float32([tr[-1] for tr in tracks]).reshape(-1, 1, 2)
		else:
			nx,ny = 20,15
			p0 = np.concatenate(np.indices((nx,ny)).swapaxes(0,2), axis=0)
			p0 = ((p0 + 0.5) * (capw / nx, caph / ny)).astype(np.float32)
			tracks = [[pt] for pt in p0]
			p0.shape = (-1, 1, 2)

		p1, st, err = cv2.calcOpticalFlowPyrLK(prevframe, curframe, p0, None, **lk_params)
		p0r, st, err = cv2.calcOpticalFlowPyrLK(curframe, prevframe, p1, None, **lk_params)
		d = abs(p0-p0r).reshape(-1, 2).max(-1)
		good = (d < 1)

		new_tracks = []
		for tr, (x,y), good_flag in zip(tracks, p1.reshape(-1, 2), good):
			if not good_flag:
				continue

			tr.append((x,y))
			if len(tr) > tracklen:
				del tr[0]

			new_tracks.append(tr)
			cv2.circle(vis, (x, y), 2, (0, 255, 0), -1)

		tracks = new_tracks
		#cv2.polylines(vis, [np.int32(tr) for tr in tracks], False, (0, 255, 0), thickness=1)

		delta = p1 - p0
		p1 = p0 + delta * 5

		p0 = p0[good].reshape((-1, 2))
		p1 = p1[good].reshape((-1, 2))
		delta = p1 - p0
		p0[:,:] = (capw/2, caph/2)
		p0 += delta
		p1 = p0

		if len(delta):
			#import pdb; pdb.set_trace()
			clusterradius = 20.0
			(centroids, labels) = clustering(clusterradius, delta)
			h = {}
			for i in labels:
				if i not in h: h[i] = 0
				h[i] += 1

			best = max(h, key=lambda k: h[k])
			(bestx, besty) = np.int32(centroids[best] + (capw/2, caph/2))

			#p1 = p0 + delta * 1
			#import pdb; pdb.set_trace()
			cv2.polylines(vis, np.concatenate([p0.reshape((-1, 1, 2)),p1.reshape((-1, 1, 2))], axis=1).astype(np.int32), False, (0, 0, 255), thickness=3)

			cv2.line(vis, (bestx - 10, besty), (bestx + 10, besty), (255, 255, 0))
			cv2.line(vis, (bestx, besty - 10), (bestx, besty + 10), (255, 255, 0))
			cv2.circle(vis, (bestx, besty), int(clusterradius), (255, 255, 0), thickness=2)

	# todo: find more tracks
	mask = np.zeros_like(curframe)
	mask[:] = 255
	for x, y in [np.int32(tr[-1]) for tr in tracks]:
		cv2.circle(mask, (x, y), 5, 0, -1)
	p = cv2.goodFeaturesToTrack(curframe, mask=mask, **feature_params)
	if p is not None:
		for x, y in np.float32(p).reshape(-1, 2):
			tracks.append([(x, y)])

	return (tracks, vis)

if __name__ == '__main__':
	source = 'V:\\Video AG\\archiv\\14ws-infin-141014\\14ws-infin-141014-dozent.mp4'
	source = 0
	# 0:12:40

	framescale = 1 #0.5
	prevframe = None
	curframe = None
	tracks = []
	tracklen = 10

	numthreads = cv2.getNumberOfCPUs()

	scalerpool = ThreadPool(processes=1)
	qscaling = deque()

	flowpool = ThreadPool(processes=1) # must be serial!
	qflow = deque()

	vid = cv2.VideoCapture(source)
	#vid.set(cv2.CAP_PROP_FPS, 15)
	if isinstance(source, str):
		vid.set(cv2.CAP_PROP_POS_MSEC, 1000 * (12*60 + 40))

	capw = vid.get(cv2.CAP_PROP_FRAME_WIDTH)
	caph = vid.get(cv2.CAP_PROP_FRAME_HEIGHT)

	cv2.namedWindow("source", cv2.WINDOW_NORMAL)
	cv2.resizeWindow("source", int(capw * framescale), int(caph * framescale))

	while True:
		#print "c1:", len(qscaling) < scalerpool._processes
		if len(qscaling) < scalerpool._processes:
			(rv, frame) = vid.read()
			task = scalerpool.apply_async(scale_frame, (frame,))
			qscaling.append(task)

		#print "c2a:", len(qscaling) > 0, len(qscaling) > 0 and qscaling[0].ready()
		#print "c2b:", len(qflow) < flowpool._processes
		while len(qscaling) > 0 and qscaling[0].ready() and len(qflow) < flowpool._processes:
			prevframe = curframe
			curframe = qscaling.popleft().get()
			if prevframe is not None and curframe is not None:
				task = flowpool.apply_async(flow_frame, (tracks, prevframe, curframe))
				qflow.append(task)

		#print "c3:", len(qflow) > 0, len(qflow) > 0 and qflow[0].ready()
		while len(qflow) > 0 and qflow[0].ready():
			(tracks, vis) = qflow.popleft().get()
			cv2.imshow("source", vis)

		key = cv2.waitKey(1)
		if key == -1: continue
		if key == 27: break

	cv2.destroyAllWindows()
<commit_msg>bgtrack: it's not that easy<commit_after>from __future__ import division

import os
import sys
import math
import numpy as np
import cv2
from multiprocessing.pool import ThreadPool
from collections import deque

lk_params = dict(
	winSize = (15,15),
	maxLevel = 2,
	criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03)
)

feature_params = dict(
	maxCorners = 500,
	qualityLevel = 0.3,
	minDistance = 9,
	blockSize = 7
)

def meanshift(p, radius, points):
	while True:
		pold = p
		inliers = (np.linalg.norm(points - p, axis=1) <= radius)
		p = np.mean(points[inliers], axis=0)
		shift = np.linalg.norm(p - pold)
		if shift < radius * 0.1:
			break

	return p

def clustering(radius, points):
	centroids = np.zeros((0,2), dtype=np.float32)
	labels = np.zeros((len(points),), dtype=np.int0)
	labels[:] = -1
	for i,pt in enumerate(points):
		pt1 = meanshift(pt, radius, points)
		# find nearest centroid
		if len(centroids) == 0:
			labels[i] = 0
			centroids = np.vstack([centroids, pt1])
		else:
			d = np.linalg.norm(centroids - pt1, axis=1)
			imin = np.argmin(d)
			if d[imin] < radius:
				labels[i] = imin
			else:
				labels[i] = len(centroids)
				centroids = np.vstack([centroids, pt1])

	return (centroids, labels)

class DummyTask:
    def __init__(self, fn, args):
        self.data = fn(*args)
    def ready(self):
        return True
    def get(self):
        return self.data

def to_f8(array):
	return tuple(np.rint(array * 256).astype('int'))

def scale_frame(frame):
	frame = cv2.resize(frame, None, fx=framescale, fy=framescale, interpolation=cv2.INTER_AREA)
	framegray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
	return framegray

def flow_frame_farneback(tracks, prevframe, curframe):

	flow = cv2.calcOpticalFlowFarneback(
		prevframe, curframe,
		pyr_scale=0.5,
		levels=7,
		winsize=5,
		iterations=3,
		poly_n=5,
		poly_sigma=1.1,
		flow=None,
		flags=0
	)

	cflow = flow.view(np.complex64).squeeze() / framescale
	angle = np.angle(cflow) % (2*np.pi)
	magnitude = np.abs(cflow)

	# HSV
	#vis = cv2.cvtColor(curframe, cv2.COLOR_GRAY2BGR)
	vis = np.zeros(flow.shape[:2] + (3,), dtype=np.uint8)

	vis[:,:,0] = np.uint8(angle / np.pi * 90)
	vis[:,:,1] = 255
	vis[:,:,2] = magnitude / (magnitude + 10) * 255

	vis = cv2.cvtColor(vis, cv2.COLOR_HSV2BGR)

	return vis

def fix8(x):
	if isinstance(x, (int, float)):
		return int(0.5 + x * (1<<8))
	else:
		return tuple(np.int32(np.array(x) * (1<<8) + 0.5))

def flow_frame(tracks, prevframe, curframe):
	assert prevframe is not None and curframe is not None

	vis = cv2.cvtColor(curframe, cv2.COLOR_GRAY2BGR)

	if len(tracks) > 0:
		if 1:
			p0 = np.float32([tr[-1] for tr in tracks]).reshape(-1, 1, 2)
		else:
			nx,ny = 20,15
			p0 = np.concatenate(np.indices((nx,ny)).swapaxes(0,2), axis=0)
			p0 = ((p0 + 0.5) * (capw / nx * framescale, caph / ny * framescale)).astype(np.float32)
			tracks = [[pt] for pt in p0]
			p0.shape = (-1, 1, 2)

		p1, st, err = cv2.calcOpticalFlowPyrLK(prevframe, curframe, p0, None, **lk_params)
		p0r, st, err = cv2.calcOpticalFlowPyrLK(curframe, prevframe, p1, None, **lk_params)
		d = abs(p0-p0r).reshape(-1, 2).max(-1)
		good = (d < 1)

		new_tracks = []
		for tr, (x,y), good_flag in zip(tracks, p1.reshape(-1, 2), good):
			if not good_flag:
				continue

			tr.append((x,y))
			if len(tr) > tracklen:
				del tr[0]

			new_tracks.append(tr)
			cv2.circle(vis, fix8([x, y]), fix8(2), (0, 255, 0), -1, shift=8, lineType=cv2.LINE_AA)

		tracks = new_tracks
		#cv2.polylines(vis, [np.int32(tr) for tr in tracks], False, (0, 255, 0), thickness=1)

		delta = p1 - p0
		p1 = p0 + delta * 5

		center = np.array([capw, caph]) * 0.5 * framescale

		p0 = p0[good].reshape((-1, 2))
		p1 = p1[good].reshape((-1, 2))
		delta = p1 - p0
		p0[:,:] = center
		p0 += delta
		p1 = p0

		if len(delta):
			#import pdb; pdb.set_trace()
			clusterradius = (caph / 20) * framescale
			(centroids, labels) = clustering(clusterradius, delta)
			h = np.bincount(labels)
			best = np.argmax(h)
			(bestx, besty) = bestp = centroids[best] + center

			#p1 = p0 + delta * 1
			#import pdb; pdb.set_trace()
			cv2.polylines(vis, fix8(np.concatenate([p0.reshape((-1, 1, 2)),p1.reshape((-1, 1, 2))], axis=1)), False, (0, 0, 255), thickness=3, shift=8, lineType=cv2.LINE_AA)

			cv2.line(vis, fix8(center - (5, 0)), fix8(center + (5, 0)), (255, 0, 0), shift=8, lineType=cv2.LINE_AA)
			cv2.line(vis, fix8(center - (0, 5)), fix8(center + (0, 5)), (255, 0, 0), shift=8, lineType=cv2.LINE_AA)

			cv2.line(vis, fix8(bestp - (10, 0)), fix8(bestp + (10, 0)), (255, 0, 255), shift=8, lineType=cv2.LINE_AA)
			cv2.line(vis, fix8(bestp - (0, 10)), fix8(bestp + (0, 10)), (255, 0, 255), shift=8, lineType=cv2.LINE_AA)

			cv2.circle(vis, fix8((bestx, besty)), fix8(clusterradius), (255, 0, 255), thickness=2, shift=8, lineType=cv2.LINE_AA)

	# todo: find more tracks
	mask = np.zeros_like(curframe)
	mask[:] = 255
	for x, y in [np.int32(tr[-1]) for tr in tracks]:
		cv2.circle(mask, (x, y), 5, 0, -1)
	p = cv2.goodFeaturesToTrack(curframe, mask=mask, **feature_params)
	if p is not None:
		for x, y in np.float32(p).reshape(-1, 2):
			tracks.append([(x, y)])

	return (tracks, vis)

if __name__ == '__main__':
	source = 'V:\\Video AG\\archiv\\14ws-infin-141014\\14ws-infin-141014-dozent.mp4'
	# 0:12:40
	source = 0

	if sys.argv[1:]:
		source = sys.argv[1]

	framescale = 0.25 #0.5
	prevframe = None
	curframe = None
	tracks = []
	tracklen = 10
	flowmethod = 0

	numthreads = cv2.getNumberOfCPUs()

	scalerpool = ThreadPool(processes=1)
	qscaling = deque()

	flowpool = ThreadPool(processes=1) # must be serial!
	qflow = deque()

	vid = cv2.VideoCapture(source)
	#vid.set(cv2.CAP_PROP_FPS, 15)
	#if isinstance(source, str):
	#	vid.set(cv2.CAP_PROP_POS_MSEC, 1000 * (12*60 + 40))

	vid.set(cv2.CAP_PROP_POS_MSEC, 1000 * (10*60 + 0))

	if source == 0:
		vid.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
		vid.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

	capw = vid.get(cv2.CAP_PROP_FRAME_WIDTH)
	caph = vid.get(cv2.CAP_PROP_FRAME_HEIGHT)

	cv2.namedWindow("source", cv2.WINDOW_NORMAL)
	cv2.resizeWindow("source", int(capw), int(caph))

	while True:
		#print "c1:", len(qscaling) < scalerpool._processes
		if len(qscaling) < scalerpool._processes:
			(rv, frame) = vid.read()
			task = scalerpool.apply_async(scale_frame, (frame,))
			qscaling.append(task)

		#print "c2a:", len(qscaling) > 0, len(qscaling) > 0 and qscaling[0].ready()
		#print "c2b:", len(qflow) < flowpool._processes
		while len(qscaling) > 0 and qscaling[0].ready() and len(qflow) < flowpool._processes:
			prevframe = curframe
			curframe = qscaling.popleft().get()
			cv2.imshow("scaled", curframe)
			if prevframe is not None and curframe is not None:
				if flowmethod == 0:
					task = flowpool.apply_async(flow_frame, (tracks, prevframe, curframe))
				elif flowmethod == 1:
					task = flowpool.apply_async(flow_frame_farneback, (tracks, prevframe, curframe))
				qflow.append(task)

		#print "c3:", len(qflow) > 0, len(qflow) > 0 and qflow[0].ready()
		while len(qflow) > 0 and qflow[0].ready():
			if flowmethod == 0:
				(tracks, vis) = qflow.popleft().get()
			elif flowmethod == 1:
				vis = qflow.popleft().get()
			cv2.imshow("source", vis)

		key = cv2.waitKey(1)
		if key == -1: continue
		if key == 27: break

	cv2.destroyAllWindows()
<|endoftext|>


<commit_before># Copyright (c) 2015, Frappe Technologies Pvt. Ltd. and Contributors
# MIT License. See license.txt

from __future__ import unicode_literals

import frappe
from frappe import _
from frappe.website.website_generator import WebsiteGenerator
from frappe.website.render import clear_cache
from frappe.utils import today, cint, global_date_format, get_fullname, strip_html_tags, markdown, sanitize_html
from frappe.website.utils import (find_first_image, get_html_content_based_on_type,
	get_comment_list)

class BlogPost(WebsiteGenerator):
	website = frappe._dict(
		order_by = "published_on desc"
	)

	def make_route(self):
		if not self.route:
			return frappe.db.get_value('Blog Category', self.blog_category,
				'route') + '/' + self.scrub(self.title)

	def get_feed(self):
		return self.title

	def validate(self):
		super(BlogPost, self).validate()

		if not self.blog_intro:
			content = get_html_content_based_on_type(self, 'content', self.content_type)
			self.blog_intro = content[:140]
			self.blog_intro = strip_html_tags(self.blog_intro)

		if self.blog_intro:
			self.blog_intro = self.blog_intro[:140]

		if self.published and not self.published_on:
			self.published_on = today()

		# update posts
		frappe.db.sql("""UPDATE `tabBlogger` SET `posts`=(SELECT COUNT(*) FROM `tabBlog Post`
			WHERE IFNULL(`blogger`,'')=`tabBlogger`.`name`)
			WHERE `name`=%s""", (self.blogger,))

	def on_update(self):
		super(BlogPost, self).on_update()
		clear_cache("writers")

	def on_trash(self):
		super(BlogPost, self).on_trash()

	def get_context(self, context):
		# this is for double precaution. usually it wont reach this code if not published
		if not cint(self.published):
			raise Exception("This blog has not been published yet!")

		# temp fields
		context.full_name = get_fullname(self.owner)
		context.updated = global_date_format(self.published_on)

		if self.blogger:
			context.blogger_info = frappe.get_doc("Blogger", self.blogger).as_dict()
			context.author = self.blogger

		context.content = get_html_content_based_on_type(self, 'content', self.content_type)

		#if meta description is not present, then blog intro or first 140 characters of the blog will be set as description
		context.description = self.meta_description or self.blog_intro or strip_html_tags(context.content[:140])

		context.metatags = {
			"name": self.title,
			"description": context.description,
		}

		#if meta image is not present, then first image inside the blog will be set as the meta image
		image = find_first_image(context.content)
		context.metatags["image"] = self.meta_image or image

		self.load_comments(context)

		context.category = frappe.db.get_value("Blog Category",
			context.doc.blog_category, ["title", "route"], as_dict=1)
		context.parents = [{"name": _("Home"), "route":"/"},
			{"name": "Blog", "route": "/blog"},
			{"label": context.category.title, "route":context.category.route}]

	def load_comments(self, context):
		context.comment_list = get_comment_list(self.doctype, self.name)

		if not context.comment_list:
			context.comment_text = _('No comments yet')
		else:
			if(len(context.comment_list)) == 1:
				context.comment_text = _('1 comment')
			else:
				context.comment_text = _('{0} comments').format(len(context.comment_list))

def get_list_context(context=None):
	list_context = frappe._dict(
		template = "templates/includes/blog/blog.html",
		get_list = get_blog_list,
		hide_filters = True,
		children = get_children(),
		# show_search = True,
		title = _('Blog')
	)

	category = sanitize_html(frappe.local.form_dict.blog_category or frappe.local.form_dict.category)
	if category:
		category_title = get_blog_category(category)
		list_context.sub_title = _("Posts filed under {0}").format(category_title)
		list_context.title = category_title

	elif frappe.local.form_dict.blogger:
		blogger = frappe.db.get_value("Blogger", {"name": frappe.local.form_dict.blogger}, "full_name")
		list_context.sub_title = _("Posts by {0}").format(blogger)
		list_context.title = blogger

	elif frappe.local.form_dict.txt:
		list_context.sub_title = _('Filtered by "{0}"').format(sanitize_html(frappe.local.form_dict.txt))

	if list_context.sub_title:
		list_context.parents = [{"name": _("Home"), "route": "/"},
								{"name": "Blog", "route": "/blog"}]
	else:
		list_context.parents = [{"name": _("Home"), "route": "/"}]

	list_context.update(frappe.get_doc("Blog Settings", "Blog Settings").as_dict(no_default_fields=True))
	return list_context

def get_children():
	return frappe.db.sql("""select route as name,
		title from `tabBlog Category`
		where published = 1
		and exists (select name from `tabBlog Post`
			where `tabBlog Post`.blog_category=`tabBlog Category`.name and published=1)
		order by title asc""", as_dict=1)

def clear_blog_cache():
	for blog in frappe.db.sql_list("""select route from
		`tabBlog Post` where ifnull(published,0)=1"""):
		clear_cache(blog)

	clear_cache("writers")

def get_blog_category(route):
	return frappe.db.get_value("Blog Category", {"name": route}, "title") or route

def get_blog_list(doctype, txt=None, filters=None, limit_start=0, limit_page_length=20, order_by=None):
	conditions = []
	category = filters.blog_category or sanitize_html(frappe.local.form_dict.blog_category or frappe.local.form_dict.category)
	if filters:
		if filters.blogger:
			conditions.append('t1.blogger=%s' % frappe.db.escape(filters.blogger))
	if category:
		conditions.append('t1.blog_category=%s' % frappe.db.escape(category))

	if txt:
		conditions.append('(t1.content like {0} or t1.title like {0}")'.format(frappe.db.escape('%' + txt + '%')))

	if conditions:
		frappe.local.no_cache = 1

	query = """\
		select
			t1.title, t1.name, t1.blog_category, t1.route, t1.published_on,
				t1.published_on as creation,
				t1.content as content,
				t1.content_type as content_type,
				t1.content_html as content_html,
				t1.content_md as content_md,
				ifnull(t1.blog_intro, t1.content) as intro,
				t2.full_name, t2.avatar, t1.blogger,
				(select count(name) from `tabComment`
					where
						comment_type='Comment'
						and reference_doctype='Blog Post'
						and reference_name=t1.name) as comments
		from `tabBlog Post` t1, `tabBlogger` t2
		where ifnull(t1.published,0)=1
		and t1.blogger = t2.name
		%(condition)s
		order by published_on desc, name asc
		limit %(start)s, %(page_len)s""" % {
			"start": limit_start, "page_len": limit_page_length,
				"condition": (" and " + " and ".join(conditions)) if conditions else ""
		}

	posts = frappe.db.sql(query, as_dict=1)

	for post in posts:

		post.content = get_html_content_based_on_type(post, 'content', post.content_type)
		post.cover_image = find_first_image(post.content)
		post.published = global_date_format(post.creation)
		post.content = strip_html_tags(post.content)

		if not post.comments:
			post.comment_text = _('No comments yet')
		elif post.comments==1:
			post.comment_text = _('1 comment')
		else:
			post.comment_text = _('{0} comments').format(str(post.comments))

		post.avatar = post.avatar or ""
		post.category = frappe.db.get_value('Blog Category', post.blog_category,
			['route', 'title'], as_dict=True)

		if post.avatar and (not "http:" in post.avatar and not "https:" in post.avatar) and not post.avatar.startswith("/"):
			post.avatar = "/" + post.avatar

	return posts
<commit_msg>refactor: explicitly set None<commit_after># Copyright (c) 2015, Frappe Technologies Pvt. Ltd. and Contributors
# MIT License. See license.txt

from __future__ import unicode_literals

import frappe
from frappe import _
from frappe.website.website_generator import WebsiteGenerator
from frappe.website.render import clear_cache
from frappe.utils import today, cint, global_date_format, get_fullname, strip_html_tags, markdown, sanitize_html
from frappe.website.utils import (find_first_image, get_html_content_based_on_type,
	get_comment_list)

class BlogPost(WebsiteGenerator):
	website = frappe._dict(
		order_by = "published_on desc"
	)

	def make_route(self):
		if not self.route:
			return frappe.db.get_value('Blog Category', self.blog_category,
				'route') + '/' + self.scrub(self.title)

	def get_feed(self):
		return self.title

	def validate(self):
		super(BlogPost, self).validate()

		if not self.blog_intro:
			content = get_html_content_based_on_type(self, 'content', self.content_type)
			self.blog_intro = content[:140]
			self.blog_intro = strip_html_tags(self.blog_intro)

		if self.blog_intro:
			self.blog_intro = self.blog_intro[:140]

		if self.published and not self.published_on:
			self.published_on = today()

		# update posts
		frappe.db.sql("""UPDATE `tabBlogger` SET `posts`=(SELECT COUNT(*) FROM `tabBlog Post`
			WHERE IFNULL(`blogger`,'')=`tabBlogger`.`name`)
			WHERE `name`=%s""", (self.blogger,))

	def on_update(self):
		super(BlogPost, self).on_update()
		clear_cache("writers")

	def on_trash(self):
		super(BlogPost, self).on_trash()

	def get_context(self, context):
		# this is for double precaution. usually it wont reach this code if not published
		if not cint(self.published):
			raise Exception("This blog has not been published yet!")

		# temp fields
		context.full_name = get_fullname(self.owner)
		context.updated = global_date_format(self.published_on)

		if self.blogger:
			context.blogger_info = frappe.get_doc("Blogger", self.blogger).as_dict()
			context.author = self.blogger

		context.content = get_html_content_based_on_type(self, 'content', self.content_type)

		#if meta description is not present, then blog intro or first 140 characters of the blog will be set as description
		context.description = self.meta_description or self.blog_intro or strip_html_tags(context.content[:140])

		context.metatags = {
			"name": self.title,
			"description": context.description,
		}

		#if meta image is not present, then first image inside the blog will be set as the meta image
		image = find_first_image(context.content)
		context.metatags["image"] = self.meta_image or image or None

		self.load_comments(context)

		context.category = frappe.db.get_value("Blog Category",
			context.doc.blog_category, ["title", "route"], as_dict=1)
		context.parents = [{"name": _("Home"), "route":"/"},
			{"name": "Blog", "route": "/blog"},
			{"label": context.category.title, "route":context.category.route}]

	def load_comments(self, context):
		context.comment_list = get_comment_list(self.doctype, self.name)

		if not context.comment_list:
			context.comment_text = _('No comments yet')
		else:
			if(len(context.comment_list)) == 1:
				context.comment_text = _('1 comment')
			else:
				context.comment_text = _('{0} comments').format(len(context.comment_list))

def get_list_context(context=None):
	list_context = frappe._dict(
		template = "templates/includes/blog/blog.html",
		get_list = get_blog_list,
		hide_filters = True,
		children = get_children(),
		# show_search = True,
		title = _('Blog')
	)

	category = sanitize_html(frappe.local.form_dict.blog_category or frappe.local.form_dict.category)
	if category:
		category_title = get_blog_category(category)
		list_context.sub_title = _("Posts filed under {0}").format(category_title)
		list_context.title = category_title

	elif frappe.local.form_dict.blogger:
		blogger = frappe.db.get_value("Blogger", {"name": frappe.local.form_dict.blogger}, "full_name")
		list_context.sub_title = _("Posts by {0}").format(blogger)
		list_context.title = blogger

	elif frappe.local.form_dict.txt:
		list_context.sub_title = _('Filtered by "{0}"').format(sanitize_html(frappe.local.form_dict.txt))

	if list_context.sub_title:
		list_context.parents = [{"name": _("Home"), "route": "/"},
								{"name": "Blog", "route": "/blog"}]
	else:
		list_context.parents = [{"name": _("Home"), "route": "/"}]

	list_context.update(frappe.get_doc("Blog Settings", "Blog Settings").as_dict(no_default_fields=True))
	return list_context

def get_children():
	return frappe.db.sql("""select route as name,
		title from `tabBlog Category`
		where published = 1
		and exists (select name from `tabBlog Post`
			where `tabBlog Post`.blog_category=`tabBlog Category`.name and published=1)
		order by title asc""", as_dict=1)

def clear_blog_cache():
	for blog in frappe.db.sql_list("""select route from
		`tabBlog Post` where ifnull(published,0)=1"""):
		clear_cache(blog)

	clear_cache("writers")

def get_blog_category(route):
	return frappe.db.get_value("Blog Category", {"name": route}, "title") or route

def get_blog_list(doctype, txt=None, filters=None, limit_start=0, limit_page_length=20, order_by=None):
	conditions = []
	category = filters.blog_category or sanitize_html(frappe.local.form_dict.blog_category or frappe.local.form_dict.category)
	if filters:
		if filters.blogger:
			conditions.append('t1.blogger=%s' % frappe.db.escape(filters.blogger))
	if category:
		conditions.append('t1.blog_category=%s' % frappe.db.escape(category))

	if txt:
		conditions.append('(t1.content like {0} or t1.title like {0}")'.format(frappe.db.escape('%' + txt + '%')))

	if conditions:
		frappe.local.no_cache = 1

	query = """\
		select
			t1.title, t1.name, t1.blog_category, t1.route, t1.published_on,
				t1.published_on as creation,
				t1.content as content,
				t1.content_type as content_type,
				t1.content_html as content_html,
				t1.content_md as content_md,
				ifnull(t1.blog_intro, t1.content) as intro,
				t2.full_name, t2.avatar, t1.blogger,
				(select count(name) from `tabComment`
					where
						comment_type='Comment'
						and reference_doctype='Blog Post'
						and reference_name=t1.name) as comments
		from `tabBlog Post` t1, `tabBlogger` t2
		where ifnull(t1.published,0)=1
		and t1.blogger = t2.name
		%(condition)s
		order by published_on desc, name asc
		limit %(start)s, %(page_len)s""" % {
			"start": limit_start, "page_len": limit_page_length,
				"condition": (" and " + " and ".join(conditions)) if conditions else ""
		}

	posts = frappe.db.sql(query, as_dict=1)

	for post in posts:

		post.content = get_html_content_based_on_type(post, 'content', post.content_type)
		post.cover_image = find_first_image(post.content)
		post.published = global_date_format(post.creation)
		post.content = strip_html_tags(post.content)

		if not post.comments:
			post.comment_text = _('No comments yet')
		elif post.comments==1:
			post.comment_text = _('1 comment')
		else:
			post.comment_text = _('{0} comments').format(str(post.comments))

		post.avatar = post.avatar or ""
		post.category = frappe.db.get_value('Blog Category', post.blog_category,
			['route', 'title'], as_dict=True)

		if post.avatar and (not "http:" in post.avatar and not "https:" in post.avatar) and not post.avatar.startswith("/"):
			post.avatar = "/" + post.avatar

	return posts
<|endoftext|>


<commit_before>from __future__ import absolute_import

import os
import subprocess

from pip._vendor.six.moves.urllib import request as urllib_request

from pip._internal.utils.misc import hide_url
from pip._internal.utils.typing import MYPY_CHECK_RUNNING
from pip._internal.vcs import bazaar, git, mercurial, subversion
from tests.lib import path_to_url

if MYPY_CHECK_RUNNING:
    from tests.lib.path import Path

def _create_initools_repository(directory):
    subprocess.check_call('svnadmin create INITools'.split(), cwd=directory)

def _dump_initools_repository(directory):
    filename, _ = urllib_request.urlretrieve(
        'http://bitbucket.org/hltbra/pip-initools-dump/raw/8b55c908a320/'
        'INITools_modified.dump'
    )
    initools_folder = os.path.join(directory, 'INITools')
    devnull = open(os.devnull, 'w')
    dump = open(filename)
    subprocess.check_call(
        ['svnadmin', 'load', initools_folder],
        stdin=dump,
        stdout=devnull,
    )
    dump.close()
    devnull.close()
    os.remove(filename)

def _create_svn_repository_for_initools(directory):
    if not os.path.exists(os.path.join(directory, 'INITools')):
        _create_initools_repository(directory)
        _dump_initools_repository(directory)

def _get_vcs_and_checkout_url(remote_repository, directory):
    vcs_classes = {'svn': subversion.Subversion,
                   'git': git.Git,
                   'bzr': bazaar.Bazaar,
                   'hg': mercurial.Mercurial}
    default_vcs = 'svn'
    if '+' not in remote_repository:
        remote_repository = '%s+%s' % (default_vcs, remote_repository)
    vcs, repository_path = remote_repository.split('+', 1)
    vcs_class = vcs_classes[vcs]
    branch = ''
    if vcs == 'svn':
        branch = os.path.basename(remote_repository)
        # remove the slash
        repository_name = os.path.basename(
            remote_repository[:-len(branch) - 1]
        )
    else:
        repository_name = os.path.basename(remote_repository)

    destination_path = os.path.join(directory, repository_name)
    if not os.path.exists(destination_path):
        url = hide_url(remote_repository)
        vcs_class().obtain(destination_path, url=url)
    return '%s+%s' % (
        vcs,
        path_to_url('/'.join([directory, repository_name, branch])),
    )

def local_checkout(
    remote_repo,  # type: str
    temp_path,    # type: Path
):
    # type: (...) -> str
    """
    :param temp_path: the return value of the tmpdir fixture, which is a
        temp directory Path object unique to each test function invocation,
        created as a sub directory of the base temp directory.
    """
    directory = temp_path.joinpath('cache')
    if not os.path.exists(directory):
        os.mkdir(directory)
        # os.makedirs(directory)

    if remote_repo.startswith('svn'):
        _create_svn_repository_for_initools(directory)
    return _get_vcs_and_checkout_url(remote_repo, directory)

def local_repo(remote_repo, temp_path):
    return local_checkout(remote_repo, temp_path).split('+', 1)[1]
<commit_msg>Simplify local_repos.py.<commit_after>from __future__ import absolute_import

import os
import subprocess

from pip._vendor.six.moves.urllib import request as urllib_request

from pip._internal.utils.misc import hide_url
from pip._internal.utils.typing import MYPY_CHECK_RUNNING
from pip._internal.vcs import vcs
from tests.lib import path_to_url

if MYPY_CHECK_RUNNING:
    from tests.lib.path import Path

def _ensure_svn_initools_repo(directory):
    """
    Create the SVN INITools repo if it doesn't already exist.
    """
    initools_dir = os.path.join(directory, 'INITools')
    repo_url_path = os.path.join(initools_dir, 'trunk')
    if os.path.exists(initools_dir):
        return repo_url_path

    subprocess.check_call('svnadmin create INITools'.split(), cwd=directory)

    filename, _ = urllib_request.urlretrieve(
        'http://bitbucket.org/hltbra/pip-initools-dump/raw/8b55c908a320/'
        'INITools_modified.dump'
    )
    devnull = open(os.devnull, 'w')
    dump = open(filename)
    subprocess.check_call(
        ['svnadmin', 'load', initools_dir],
        stdin=dump,
        stdout=devnull,
    )
    dump.close()
    devnull.close()
    os.remove(filename)

    return repo_url_path

def local_checkout(
    remote_repo,  # type: str
    temp_path,    # type: Path
):
    # type: (...) -> str
    """
    :param temp_path: the return value of the tmpdir fixture, which is a
        temp directory Path object unique to each test function invocation,
        created as a sub directory of the base temp directory.
    """
    assert '+' in remote_repo
    vcs_name, repository_path = remote_repo.split('+', 1)

    directory = temp_path.joinpath('cache')
    if not os.path.exists(directory):
        os.mkdir(directory)

    if vcs_name == 'svn':
        assert remote_repo.endswith('/INITools/trunk')
        repo_url_path = _ensure_svn_initools_repo(directory)
    else:
        repository_name = os.path.basename(remote_repo)
        repo_url_path = os.path.join(directory, repository_name)
        vcs_backend = vcs.get_backend(vcs_name)
        vcs_backend.obtain(repo_url_path, url=hide_url(remote_repo))

    return '{}+{}'.format(vcs_name, path_to_url(repo_url_path))

def local_repo(remote_repo, temp_path):
    return local_checkout(remote_repo, temp_path).split('+', 1)[1]
<|endoftext|>


<commit_before>import pytest
import vaex
pytest.importorskip("sklearn")
from vaex.ml.sklearn import SKLearnPredictor, IncrementalPredictor

import numpy as np

# Regressions
from sklearn.linear_model import LinearRegression, Ridge, Lasso, SGDClassifier, SGDRegressor
from sklearn.svm import SVR
from sklearn.ensemble import AdaBoostRegressor, GradientBoostingRegressor, RandomForestRegressor

# Classifiers
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from sklearn.ensemble import AdaBoostClassifier, GradientBoostingClassifier, RandomForestClassifier

models_regression = [LinearRegression(),
                     Ridge(random_state=42, max_iter=100),
                     Lasso(random_state=42, max_iter=100),
                     SVR(gamma='scale'),
                     AdaBoostRegressor(random_state=42, n_estimators=10),
                     GradientBoostingRegressor(random_state=42, max_depth=3, n_estimators=10),
                     RandomForestRegressor(n_estimators=10, random_state=42, max_depth=3)]

models_classification = [LogisticRegression(solver='lbfgs', max_iter=100, random_state=42),
                         SVC(gamma='scale', max_iter=100),
                         AdaBoostClassifier(random_state=42, n_estimators=10),
                         GradientBoostingClassifier(random_state=42, max_depth=3, n_estimators=10),
                         RandomForestClassifier(n_estimators=10, random_state=42, max_depth=3)]

def test_sklearn_estimator():
    ds = vaex.ml.datasets.load_iris()
    features = ['sepal_length', 'sepal_width', 'petal_length']

    train, test = ds.ml.train_test_split(verbose=False)

    model = SKLearnPredictor(model=LinearRegression(), features=features, prediction_name='pred')
    model.fit(train, train.petal_width)
    prediction = model.predict(test)
    test = model.transform(test)
    np.testing.assert_array_almost_equal(test.pred.values, prediction, decimal=5)

    # Transfer the state of train to ds
    train = model.transform(train)
    state = train.state_get()
    ds.state_set(state)
    assert ds.pred.values.shape == (150,)

def test_sklearn_estimator_virtual_columns():
    ds = vaex.ml.datasets.load_iris()
    ds['x'] = ds.sepal_length * 1
    ds['y'] = ds.sepal_width * 1
    ds['w'] = ds.petal_length * 1
    ds['z'] = ds.petal_width * 1
    train, test = ds.ml.train_test_split(test_size=0.2, verbose=False)
    features = ['x', 'y', 'z']
    model = SKLearnPredictor(model=LinearRegression(), features=features, prediction_name='pred')
    model.fit(ds, ds.w)
    ds = model.transform(ds)
    assert ds.pred.values.shape == (150,)

def test_sklearn_estimator_serialize(tmpdir):
    ds = vaex.ml.datasets.load_iris()
    features = ['sepal_length', 'sepal_width', 'petal_length']

    model = SKLearnPredictor(model=LinearRegression(), features=features, prediction_name='pred')
    model.fit(ds, ds.petal_width)

    pipeline = vaex.ml.Pipeline([model])
    pipeline.save(str(tmpdir.join('test.json')))
    pipeline.load(str(tmpdir.join('test.json')))

    model = SKLearnPredictor(model=LinearRegression(), features=features, prediction_name='pred')
    model.fit(ds, ds.petal_width)

    model.state_set(model.state_get())
    pipeline = vaex.ml.Pipeline([model])
    pipeline.save(str(tmpdir.join('test.json')))
    pipeline.load(str(tmpdir.join('test.json')))

def test_sklearn_estimator_regression_validation():
    ds = vaex.ml.datasets.load_iris()
    train, test = ds.ml.train_test_split(verbose=False)
    features = ['sepal_length', 'sepal_width', 'petal_length']

    # Dense features
    Xtrain = train[features].values
    Xtest = test[features].values
    ytrain = train.petal_width.values

    for model in models_regression:

        # vaex
        vaex_model = SKLearnPredictor(model=model, features=features, prediction_name='pred')
        vaex_model.fit(train, train.petal_width)
        test = vaex_model.transform(test)

        # sklearn
        model.fit(Xtrain, ytrain)
        skl_pred = model.predict(Xtest)

        np.testing.assert_array_almost_equal(test.pred.values, skl_pred, decimal=5)

def test_sklearn_estimator_pipeline():
    ds = vaex.ml.datasets.load_iris()
    train, test = ds.ml.train_test_split(verbose=False)
    # Add virtual columns
    train['sepal_virtual'] = np.sqrt(train.sepal_length**2 + train.sepal_width**2)
    train['petal_scaled'] = train.petal_length * 0.2
    # Do a pca
    features = ['sepal_virtual', 'petal_scaled']
    pca = train.ml.pca(n_components=2, features=features)
    train = pca.transform(train)
    # Do state transfer
    st = train.ml.state_transfer()
    # now apply the model
    features = ['sepal_virtual', 'petal_scaled']
    model = SKLearnPredictor(model=LinearRegression(), features=features, prediction_name='pred')
    model.fit(train, train.petal_width)
    # Create a pipeline
    pipeline = vaex.ml.Pipeline([st, model])
    # Use the pipeline
    pred = pipeline.predict(test)
    df_trans = pipeline.transform(test)

    # WARNING: on windows/appveyor this gives slightly different results
    # do we fully understand why? I also have the same results on my osx laptop
    # sklearn 0.21.1 (scikit-learn-0.21.2 is installed on windows) so it might be a
    # version related thing
    np.testing.assert_array_almost_equal(pred, df_trans.pred.values)

def test_sklearn_estimator_classification_validation():
    ds = vaex.ml.datasets.load_titanic()

    train, test = ds.ml.train_test_split(verbose=False)
    features = ['pclass', 'parch', 'sibsp']

    # Dense features
    Xtrain = train[features].values
    Xtest = test[features].values
    ytrain = train.survived.values

    for model in models_classification:

        # vaex
        vaex_model = SKLearnPredictor(model=model, features=features, prediction_name='pred')
        vaex_model.fit(train, train.survived)
        test = vaex_model.transform(test)

        # scikit-learn
        model.fit(Xtrain, ytrain)
        skl_pred = model.predict(Xtest)

        assert np.all(skl_pred == test.pred.values)

def test_sklearn_incremental_predictor():
    df = vaex.example()
    df_train, df_test = df.ml.train_test_split(test_size=0.1, verbose=False)

    features = df_train.column_names[:6]
    target = 'FeH'

    incremental = IncrementalPredictor(model=SGDRegressor(),
                                       features=features,
                                       batch_size=10_000,
                                       num_epochs=5,
                                       shuffle=True,
                                       prediction_name='pred')
    incremental.fit(df=df_train, target=target)
    df_train = incremental.transform(df_train)

    # State transfer
    state = df_train.state_get()
    df_test.state_set(state)

    assert df_train.column_count() == df_test.column_count()
    assert df_test.pred.values.shape == (33000,)

    pred_in_memory = incremental.predict(df_test)
    np.testing.assert_array_almost_equal(pred_in_memory, df_test.pred.values, decimal=1)

def test_sklearn_incremental_predictor_serialize(tmpdir):
    df = vaex.example()
    df_train, df_test = df.ml.train_test_split(test_size=0.1, verbose=False)

    features = df_train.column_names[:6]
    target = 'FeH'

    incremental = IncrementalPredictor(model=SGDRegressor(),
                                       features=features,
                                       batch_size=10_000,
                                       num_epochs=5,
                                       shuffle=True,
                                       prediction_name='pred')
    incremental.fit(df=df_train, target=target)
    df_train = incremental.transform(df_train)

    # State transfer - serialization
    df_train.state_write(str(tmpdir.join('test.json')))
    df_test.state_load(str(tmpdir.join('test.json')))

    assert df_train.column_count() == df_test.column_count()
    assert df_test.pred.values.shape == (33000,)

    pred_in_memory = incremental.predict(df_test)
    np.testing.assert_array_almost_equal(pred_in_memory, df_test.pred.values, decimal=1)

@pytest.mark.parametrize("batch_size", [6789, 10_000])
@pytest.mark.parametrize("num_epochs", [1, 5])
def test_sklearn_incremental_predictor_partial_fit_calls(batch_size, num_epochs):
    df = vaex.example()
    df_train, df_test = df.ml.train_test_split(test_size=0.1, verbose=False)

    features = df_train.column_names[:6]
    target = 'FeH'

    N_total = len(df_train)
    num_batches = (N_total + batch_size - 1) // batch_size

    # Create a mock model for counting the number of samples seen and partial_fit calls
    class MockModel():
        def __init__(self):
            self.n_samples_ = 0
            self.n_partial_fit_calls_ = 0
        def partial_fit(self, X, y):
            self.n_samples_ += X.shape[0]
            self.n_partial_fit_calls_ += 1

    incremental = IncrementalPredictor(model=MockModel(),
                                       features=features,
                                       batch_size=batch_size,
                                       num_epochs=num_epochs,
                                       shuffle=False,
                                       prediction_name='pred')

    incremental.fit(df=df_train, target=target)
    assert incremental.model.n_samples_ == N_total * num_epochs
    assert incremental.model.n_partial_fit_calls_ == num_batches * num_epochs
<commit_msg>test(ml): make tests run on py27<commit_after>import pytest
import vaex
pytest.importorskip("sklearn")
from vaex.ml.sklearn import SKLearnPredictor, IncrementalPredictor

import numpy as np

# Regressions
from sklearn.linear_model import LinearRegression, Ridge, Lasso, SGDClassifier, SGDRegressor
from sklearn.svm import SVR
from sklearn.ensemble import AdaBoostRegressor, GradientBoostingRegressor, RandomForestRegressor

# Classifiers
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from sklearn.ensemble import AdaBoostClassifier, GradientBoostingClassifier, RandomForestClassifier

models_regression = [LinearRegression(),
                     Ridge(random_state=42, max_iter=100),
                     Lasso(random_state=42, max_iter=100),
                     SVR(gamma='scale'),
                     AdaBoostRegressor(random_state=42, n_estimators=10),
                     GradientBoostingRegressor(random_state=42, max_depth=3, n_estimators=10),
                     RandomForestRegressor(n_estimators=10, random_state=42, max_depth=3)]

models_classification = [LogisticRegression(solver='lbfgs', max_iter=100, random_state=42),
                         SVC(gamma='scale', max_iter=100),
                         AdaBoostClassifier(random_state=42, n_estimators=10),
                         GradientBoostingClassifier(random_state=42, max_depth=3, n_estimators=10),
                         RandomForestClassifier(n_estimators=10, random_state=42, max_depth=3)]

def test_sklearn_estimator():
    ds = vaex.ml.datasets.load_iris()
    features = ['sepal_length', 'sepal_width', 'petal_length']

    train, test = ds.ml.train_test_split(verbose=False)

    model = SKLearnPredictor(model=LinearRegression(), features=features, prediction_name='pred')
    model.fit(train, train.petal_width)
    prediction = model.predict(test)
    test = model.transform(test)
    np.testing.assert_array_almost_equal(test.pred.values, prediction, decimal=5)

    # Transfer the state of train to ds
    train = model.transform(train)
    state = train.state_get()
    ds.state_set(state)
    assert ds.pred.values.shape == (150,)

def test_sklearn_estimator_virtual_columns():
    ds = vaex.ml.datasets.load_iris()
    ds['x'] = ds.sepal_length * 1
    ds['y'] = ds.sepal_width * 1
    ds['w'] = ds.petal_length * 1
    ds['z'] = ds.petal_width * 1
    train, test = ds.ml.train_test_split(test_size=0.2, verbose=False)
    features = ['x', 'y', 'z']
    model = SKLearnPredictor(model=LinearRegression(), features=features, prediction_name='pred')
    model.fit(ds, ds.w)
    ds = model.transform(ds)
    assert ds.pred.values.shape == (150,)

def test_sklearn_estimator_serialize(tmpdir):
    ds = vaex.ml.datasets.load_iris()
    features = ['sepal_length', 'sepal_width', 'petal_length']

    model = SKLearnPredictor(model=LinearRegression(), features=features, prediction_name='pred')
    model.fit(ds, ds.petal_width)

    pipeline = vaex.ml.Pipeline([model])
    pipeline.save(str(tmpdir.join('test.json')))
    pipeline.load(str(tmpdir.join('test.json')))

    model = SKLearnPredictor(model=LinearRegression(), features=features, prediction_name='pred')
    model.fit(ds, ds.petal_width)

    model.state_set(model.state_get())
    pipeline = vaex.ml.Pipeline([model])
    pipeline.save(str(tmpdir.join('test.json')))
    pipeline.load(str(tmpdir.join('test.json')))

def test_sklearn_estimator_regression_validation():
    ds = vaex.ml.datasets.load_iris()
    train, test = ds.ml.train_test_split(verbose=False)
    features = ['sepal_length', 'sepal_width', 'petal_length']

    # Dense features
    Xtrain = train[features].values
    Xtest = test[features].values
    ytrain = train.petal_width.values

    for model in models_regression:

        # vaex
        vaex_model = SKLearnPredictor(model=model, features=features, prediction_name='pred')
        vaex_model.fit(train, train.petal_width)
        test = vaex_model.transform(test)

        # sklearn
        model.fit(Xtrain, ytrain)
        skl_pred = model.predict(Xtest)

        np.testing.assert_array_almost_equal(test.pred.values, skl_pred, decimal=5)

def test_sklearn_estimator_pipeline():
    ds = vaex.ml.datasets.load_iris()
    train, test = ds.ml.train_test_split(verbose=False)
    # Add virtual columns
    train['sepal_virtual'] = np.sqrt(train.sepal_length**2 + train.sepal_width**2)
    train['petal_scaled'] = train.petal_length * 0.2
    # Do a pca
    features = ['sepal_virtual', 'petal_scaled']
    pca = train.ml.pca(n_components=2, features=features)
    train = pca.transform(train)
    # Do state transfer
    st = train.ml.state_transfer()
    # now apply the model
    features = ['sepal_virtual', 'petal_scaled']
    model = SKLearnPredictor(model=LinearRegression(), features=features, prediction_name='pred')
    model.fit(train, train.petal_width)
    # Create a pipeline
    pipeline = vaex.ml.Pipeline([st, model])
    # Use the pipeline
    pred = pipeline.predict(test)
    df_trans = pipeline.transform(test)

    # WARNING: on windows/appveyor this gives slightly different results
    # do we fully understand why? I also have the same results on my osx laptop
    # sklearn 0.21.1 (scikit-learn-0.21.2 is installed on windows) so it might be a
    # version related thing
    np.testing.assert_array_almost_equal(pred, df_trans.pred.values)

def test_sklearn_estimator_classification_validation():
    ds = vaex.ml.datasets.load_titanic()

    train, test = ds.ml.train_test_split(verbose=False)
    features = ['pclass', 'parch', 'sibsp']

    # Dense features
    Xtrain = train[features].values
    Xtest = test[features].values
    ytrain = train.survived.values

    for model in models_classification:

        # vaex
        vaex_model = SKLearnPredictor(model=model, features=features, prediction_name='pred')
        vaex_model.fit(train, train.survived)
        test = vaex_model.transform(test)

        # scikit-learn
        model.fit(Xtrain, ytrain)
        skl_pred = model.predict(Xtest)

        assert np.all(skl_pred == test.pred.values)

def test_sklearn_incremental_predictor():
    df = vaex.example()
    df_train, df_test = df.ml.train_test_split(test_size=0.1, verbose=False)

    features = df_train.column_names[:6]
    target = 'FeH'

    incremental = IncrementalPredictor(model=SGDRegressor(),
                                       features=features,
                                       batch_size=10*1000,
                                       num_epochs=5,
                                       shuffle=True,
                                       prediction_name='pred')
    incremental.fit(df=df_train, target=target)
    df_train = incremental.transform(df_train)

    # State transfer
    state = df_train.state_get()
    df_test.state_set(state)

    assert df_train.column_count() == df_test.column_count()
    assert df_test.pred.values.shape == (33000,)

    pred_in_memory = incremental.predict(df_test)
    np.testing.assert_array_almost_equal(pred_in_memory, df_test.pred.values, decimal=1)

def test_sklearn_incremental_predictor_serialize(tmpdir):
    df = vaex.example()
    df_train, df_test = df.ml.train_test_split(test_size=0.1, verbose=False)

    features = df_train.column_names[:6]
    target = 'FeH'

    incremental = IncrementalPredictor(model=SGDRegressor(),
                                       features=features,
                                       batch_size=10*1000,
                                       num_epochs=5,
                                       shuffle=True,
                                       prediction_name='pred')
    incremental.fit(df=df_train, target=target)
    df_train = incremental.transform(df_train)

    # State transfer - serialization
    df_train.state_write(str(tmpdir.join('test.json')))
    df_test.state_load(str(tmpdir.join('test.json')))

    assert df_train.column_count() == df_test.column_count()
    assert df_test.pred.values.shape == (33000,)

    pred_in_memory = incremental.predict(df_test)
    np.testing.assert_array_almost_equal(pred_in_memory, df_test.pred.values, decimal=1)

@pytest.mark.parametrize("batch_size", [6789, 10*1000])
@pytest.mark.parametrize("num_epochs", [1, 5])
def test_sklearn_incremental_predictor_partial_fit_calls(batch_size, num_epochs):
    df = vaex.example()
    df_train, df_test = df.ml.train_test_split(test_size=0.1, verbose=False)

    features = df_train.column_names[:6]
    target = 'FeH'

    N_total = len(df_train)
    num_batches = (N_total + batch_size - 1) // batch_size

    # Create a mock model for counting the number of samples seen and partial_fit calls
    class MockModel():
        def __init__(self):
            self.n_samples_ = 0
            self.n_partial_fit_calls_ = 0
        def partial_fit(self, X, y):
            self.n_samples_ += X.shape[0]
            self.n_partial_fit_calls_ += 1

    incremental = IncrementalPredictor(model=MockModel(),
                                       features=features,
                                       batch_size=batch_size,
                                       num_epochs=num_epochs,
                                       shuffle=False,
                                       prediction_name='pred')

    incremental.fit(df=df_train, target=target)
    assert incremental.model.n_samples_ == N_total * num_epochs
    assert incremental.model.n_partial_fit_calls_ == num_batches * num_epochs
<|endoftext|>