<commit_before>import re
import struct
import typing
from configparser import ConfigParser
from datetime import datetime, timedelta
from enum import IntEnum
from io import BytesIO, StringIO
from typing import *

from pros.common import *
from pros.serial import bytes_to_str
from pros.serial import decode_bytes_to_str
from pros.serial.ports import list_all_comports, BasePort
from .comm_error import VEXCommError
from .crc import CRC
from .message import Message
from .vex_device import VEXDevice
from ..system_device import SystemDevice

int_str = Union[int, str]


def find_v5_ports(p_type: str):
    def filter_vex_ports(p):
        return p.vid is not None and p.vid in [0x2888, 0x0501] or \
               p.name is not None and ('VEX' in p.name or 'V5' in p.name)

    def filter_v5_ports(p, locations, names):
        return (p.location is not None and any([p.location.endswith(l) for l in locations])) or \
               (p.name is not None and any([n in p.name for n in names])) or \
               (p.description is not None and any([n in p.description for n in names]))

    ports = [p for p in list_all_comports() if filter_vex_ports(p)]

    # Initially try filtering based off of location or the name of the device.
    # Doesn't work on macOS or Jonathan's Dell, so we have a fallback (below)
    user_ports = [p for p in ports if filter_v5_ports(p, ['2'], ['User'])]
    system_ports = [p for p in ports if filter_v5_ports(p, ['0'], ['System', 'Communications'])]
    joystick_ports = []  # joystick comms are very slow/unusable
    # joystick_ports = [p for p in ports if filter_v5_ports(p, ['1'], ['Controller'])]

    # TODO: test this code path (hard)
    if len(user_ports) != len(system_ports):
        if len(user_ports) > len(system_ports):
            user_ports = [p for p in user_ports if p not in system_ports]
        else:
            system_ports = [p for p in system_ports if p not in user_ports]

    if len(user_ports) == len(system_ports) and len(user_ports) > 0:
        if p_type.lower() == 'user':
            return user_ports
        elif p_type.lower() == 'system':
            return system_ports
        else:
            raise ValueError(f'Invalid port type specified: {p_type}')

    # None of the typical filters worked, so if there are only two ports, then the lower one is always*
    # the USER? port (*always = I haven't found a guarantee)
    if len(ports) == 2:
        # natural sort based on: https://stackoverflow.com/a/16090640
        def natural_key(chunk: str):
            return [int(text) if text.isdigit() else text.lower() for text in re.split(r'(\d+)', chunk)]

        ports = sorted(ports, key=lambda p: natural_key(p.device))
        if p_type.lower() == 'user':
            return [ports[1]]
        elif p_type.lower() == 'system':
            return [ports[0]]
        else:
            raise ValueError(f'Invalid port type specified: {p_type}')
    if len(joystick_ports) > 0:
        return joystick_ports
    return []


class V5Device(VEXDevice, SystemDevice):
    vid_map = {'user': 1, 'system': 15}  # type: Dict[str, int]

    class FTCompleteOptions(IntEnum):
        DONT_RUN = 0
        RUN_IMMEDIATELY = 0b01
        RUN_SCREEN = 0b11

    VEX_CRC16 = CRC(16, 0x1021)  # CRC-16-CCIT
    VEX_CRC32 = CRC(32, 0x04C11DB7)  # CRC-32 (the one used everywhere but has no name)

    def __init__(self, port: BasePort):
        self._status = None
        super().__init__(port)

    @property
    def status(self):
        if not self._status:
            self._status = self.get_system_status()
        return self._status

    def write_program(self, file: typing.BinaryIO, remote_name: str = None, ini: ConfigParser = None, slot: int = 0,
                      file_len: int = -1, run_after: FTCompleteOptions = FTCompleteOptions.DONT_RUN,
                      target: str = 'flash', quirk: int = 0, **kwargs):
        remote_base = f'slot_{slot + 1}'
        if target == 'ddr':
            self.write_file(file, f'{remote_base}.bin', file_len=file_len, type='bin',
                            target='ddr', run_after=run_after, **kwargs)
            return
        if not isinstance(ini, ConfigParser):
            ini = ConfigParser()
        if not remote_name:
            remote_name = file.name
        if len(remote_name) > 20:
            logger(__name__).info('Truncating remote name to {} for length.'.format(remote_name[:20]))
            remote_name = remote_name[:20]
        project_ini = ConfigParser()
        from semantic_version import Spec
        default_icon = 'USER902x.bmp' if Spec('>=1.0.0-22').match(self.status['cpu0_version']) else 'USER999x.bmp'
        project_ini['program'] = {
            'version': kwargs.get('version', '0.0.0') or '0.0.0',
            'name': remote_name,
            'slot': slot,
            'icon': kwargs.get('icon', default_icon) or default_icon,
            'description': 'Created with PROS',
            'date': datetime.now().isoformat()
        }
        project_ini.update(ini)
        if (quirk & 0xff) == 1:
            # WRITE BIN FILE
            self.write_file(file, f'{remote_base}.bin', file_len=file_len, type='bin', run_after=run_after, **kwargs)
            with StringIO() as ini_str:
                project_ini.write(ini_str)
                logger(__name__).info(f'Created ini: {ini_str.getvalue()}')
                with BytesIO(ini_str.getvalue().encode(encoding='ascii')) as ini_bin:
                    # WRITE INI FILE
                    self.write_file(ini_bin, f'{remote_base}.ini', type='ini', **kwargs)
        elif (quirk & 0xff) == 0:
            # STOP PROGRAM
            self.execute_program_file('', run=False)
            with StringIO() as ini_str:
                project_ini.write(ini_str)
                logger(__name__).info(f'Created ini: {ini_str.getvalue()}')
                with BytesIO(ini_str.getvalue().encode(encoding='ascii')) as ini_bin:
                    # WRITE INI FILE
                    self.write_file(ini_bin, f'{remote_base}.ini', type='ini', **kwargs)
            # WRITE BIN FILE
            self.write_file(file, f'{remote_base}.bin', file_len=file_len, type='bin', run_after=run_after, **kwargs)
        else:
            raise ValueError(f'Unknown quirk option: {quirk}')

    def read_file(self, file: typing.IO[bytes], remote_file: str, vid: int_str = 'user', target: int_str = 'flash'):
        if isinstance(vid, str):
            vid = self.vid_map[vid.lower()]
        metadata = self.get_file_metadata_by_name(remote_file, vid=vid)
        ft_meta = self.ft_initialize(remote_file, function='download', vid=vid, target=target)
        for i in range(0, ft_meta['file_size'], ft_meta['max_packet_size']):
            packet_size = ft_meta['max_packet_size']
            if i + ft_meta['max_packet_size'] > ft_meta['file_size']:
                packet_size = ft_meta['file_size'] - i
            file.write(self.ft_read(metadata['addr'] + i, packet_size))
        self.ft_complete()

    def write_file(self, file: typing.BinaryIO, remote_file: str, file_len: int = -1,
                   run_after: FTCompleteOptions = FTCompleteOptions.DONT_RUN, **kwargs):
        if file_len < 0:
            file_len = file.seek(0, 2)
            file.seek(0, 0)
        crc32 = self.VEX_CRC32.compute(file.read(file_len))
        file.seek(-file_len, 2)
        addr = kwargs.get('addr', 0x03800000)
        logger(__name__).info('Transferring {} ({} bytes) to the V5 from {}'.format(remote_file, file_len, file))<commit_msg>Refactor INI file generation into its own function<commit_after>import re
import struct
import typing
from configparser import ConfigParser
from datetime import datetime, timedelta
from enum import IntEnum
from io import BytesIO, StringIO
from typing import *

from pros.common import *
from pros.serial import bytes_to_str
from pros.serial import decode_bytes_to_str
from pros.serial.ports import list_all_comports, BasePort
from .comm_error import VEXCommError
from .crc import CRC
from .message import Message
from .vex_device import VEXDevice
from ..system_device import SystemDevice

int_str = Union[int, str]


def find_v5_ports(p_type: str):
    def filter_vex_ports(p):
        return p.vid is not None and p.vid in [0x2888, 0x0501] or \
               p.name is not None and ('VEX' in p.name or 'V5' in p.name)

    def filter_v5_ports(p, locations, names):
        return (p.location is not None and any([p.location.endswith(l) for l in locations])) or \
               (p.name is not None and any([n in p.name for n in names])) or \
               (p.description is not None and any([n in p.description for n in names]))

    ports = [p for p in list_all_comports() if filter_vex_ports(p)]

    # Initially try filtering based off of location or the name of the device.
    # Doesn't work on macOS or Jonathan's Dell, so we have a fallback (below)
    user_ports = [p for p in ports if filter_v5_ports(p, ['2'], ['User'])]
    system_ports = [p for p in ports if filter_v5_ports(p, ['0'], ['System', 'Communications'])]
    joystick_ports = []  # joystick comms are very slow/unusable
    # joystick_ports = [p for p in ports if filter_v5_ports(p, ['1'], ['Controller'])]

    # TODO: test this code path (hard)
    if len(user_ports) != len(system_ports):
        if len(user_ports) > len(system_ports):
            user_ports = [p for p in user_ports if p not in system_ports]
        else:
            system_ports = [p for p in system_ports if p not in user_ports]

    if len(user_ports) == len(system_ports) and len(user_ports) > 0:
        if p_type.lower() == 'user':
            return user_ports
        elif p_type.lower() == 'system':
            return system_ports
        else:
            raise ValueError(f'Invalid port type specified: {p_type}')

    # None of the typical filters worked, so if there are only two ports, then the lower one is always*
    # the USER? port (*always = I haven't found a guarantee)
    if len(ports) == 2:
        # natural sort based on: https://stackoverflow.com/a/16090640
        def natural_key(chunk: str):
            return [int(text) if text.isdigit() else text.lower() for text in re.split(r'(\d+)', chunk)]

        ports = sorted(ports, key=lambda p: natural_key(p.device))
        if p_type.lower() == 'user':
            return [ports[1]]
        elif p_type.lower() == 'system':
            return [ports[0]]
        else:
            raise ValueError(f'Invalid port type specified: {p_type}')
    if len(joystick_ports) > 0:
        return joystick_ports
    return []


class V5Device(VEXDevice, SystemDevice):
    vid_map = {'user': 1, 'system': 15}  # type: Dict[str, int]

    class FTCompleteOptions(IntEnum):
        DONT_RUN = 0
        RUN_IMMEDIATELY = 0b01
        RUN_SCREEN = 0b11

    VEX_CRC16 = CRC(16, 0x1021)  # CRC-16-CCIT
    VEX_CRC32 = CRC(32, 0x04C11DB7)  # CRC-32 (the one used everywhere but has no name)

    def __init__(self, port: BasePort):
        self._status = None
        super().__init__(port)

    @property
    def status(self):
        if not self._status:
            self._status = self.get_system_status()
        return self._status

    def generate_ini_file(self, remote_name: str = None, slot: int = 0, ini: ConfigParser = None, **kwargs):
        project_ini = ConfigParser()
        from semantic_version import Spec
        default_icon = 'USER902x.bmp' if Spec('>=1.0.0-22').match(self.status['cpu0_version']) else 'USER999x.bmp'
        project_ini['program'] = {
            'version': kwargs.get('version', '0.0.0') or '0.0.0',
            'name': remote_name,
            'slot': slot,
            'icon': kwargs.get('icon', default_icon) or default_icon,
            'description': 'Created with PROS',
            'date': datetime.now().isoformat()
        }
        if ini:
            project_ini.update(ini)
        with StringIO() as ini_str:
            project_ini.write(ini_str)
            logger(__name__).info(f'Created ini: {ini_str.getvalue()}')
            return ini_str.getvalue()

    def write_program(self, file: typing.BinaryIO, remote_name: str = None, ini: ConfigParser = None, slot: int = 0,
                      file_len: int = -1, run_after: FTCompleteOptions = FTCompleteOptions.DONT_RUN,
                      target: str = 'flash', quirk: int = 0, **kwargs):
        remote_base = f'slot_{slot + 1}'
        if target == 'ddr':
            self.write_file(file, f'{remote_base}.bin', file_len=file_len, type='bin',
                            target='ddr', run_after=run_after, **kwargs)
            return
        if not isinstance(ini, ConfigParser):
            ini = ConfigParser()
        if not remote_name:
            remote_name = file.name
        if len(remote_name) > 20:
            logger(__name__).info('Truncating remote name to {} for length.'.format(remote_name[:20]))
            remote_name = remote_name[:20]

        ini_file = self.generate_ini_file(remote_name=remote_name, slot=slot, ini=ini, **kwargs)
        logger(__name__).info(f'Created ini: {ini_file}')

        if (quirk & 0xff) == 1:
            # WRITE BIN FILE
            self.write_file(file, f'{remote_base}.bin', file_len=file_len, type='bin', run_after=run_after, **kwargs)
            with BytesIO(ini_file.encode(encoding='ascii')) as ini_bin:
                # WRITE INI FILE
                self.write_file(ini_bin, f'{remote_base}.ini', type='ini', **kwargs)
        elif (quirk & 0xff) == 0:
            # STOP PROGRAM
            self.execute_program_file('', run=False)
            with BytesIO(ini_file.encode(encoding='ascii')) as ini_bin:
                # WRITE INI FILE
                self.write_file(ini_bin, f'{remote_base}.ini', type='ini', **kwargs)
            # WRITE BIN FILE
            self.write_file(file, f'{remote_base}.bin', file_len=file_len, type='bin', run_after=run_after, **kwargs)
        else:
            raise ValueError(f'Unknown quirk option: {quirk}')

    def read_file(self, file: typing.IO[bytes], remote_file: str, vid: int_str = 'user', target: int_str = 'flash'):
        if isinstance(vid, str):
            vid = self.vid_map[vid.lower()]
        metadata = self.get_file_metadata_by_name(remote_file, vid=vid)
        ft_meta = self.ft_initialize(remote_file, function='download', vid=vid, target=target)
        for i in range(0, ft_meta['file_size'], ft_meta['max_packet_size']):
            packet_size = ft_meta['max_packet_size']
            if i + ft_meta['max_packet_size'] > ft_meta['file_size']:
                packet_size = ft_meta['file_size'] - i
            file.write(self.ft_read(metadata['addr'] + i, packet_size))
        self.ft_complete()

    def write_file(self, file: typing.BinaryIO, remote_file: str, file_len: int = -1,
                   run_after: FTCompleteOptions = FTCompleteOptions.DONT_RUN, **kwargs):
        if file_len < 0:
            file_len = file.seek(0, 2)
            file.seek(0, 0)
        crc32 = self.VEX_CRC32.compute(file.read(file_len))
        file.seek(-file_len, 2)
        addr = kwargs.get('addr', 0x03800000)
        logger(__name__).info('Transferring {} ({} bytes) to the V5 from {}'.format(remote_file, file_len, file))
{'fraction_non_alphanumeric': 0.1182350233445124, 'fraction_numerical': 0.01584796475307424, 'mean_word_length': 2.5944221224296857, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 15, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4024335', 'n_tokens_mistral': 4975, 'n_tokens_neox': 4731, 'n_words': 1312}