Source code for spherov2.commands.drive

import struct
from enum import IntFlag, IntEnum

from spherov2.commands import Commands
from spherov2.helper import to_bytes
from spherov2.listeners.drive import MotorStall


class DriveFlags(IntFlag):
    FORWARD = 0x0  # 0b0
    BACKWARD = 0x1  # 0b1
    TURBO = 0x2  # 0b10
    FAST_TURN = 0x4  # 0b100
    LEFT_DIRECTION = 0x8  # 0b1000
    RIGHT_DIRECTION = 0x10  # 0b10000
    ENABLE_DRIFT = 0x20  # 0b100000


class RCDriveFlags(IntFlag):
    SLEW_LINEAR_VELOCITY = 0x1  # 0b1


class XYPositionDriveFlags(IntFlag):
    FORCE_REVERSE = 0x1  # 0b1
    AUTO_REVERSE = 0x2  # 0b10


class StabilizationIndexes(IntEnum):
    NO_CONTROL_SYSTEM = 0
    FULL_CONTROL_SYSTEM = 1
    PITCH_CONTROL_SYSTEM = 2
    ROLL_CONTROL_SYSTEM = 3
    YAW_CONTROL_SYSTEM = 4
    SPEED_AND_YAW_CONTROL_SYSTEM = 5


class RawMotorModes(IntEnum):
    OFF = 0
    FORWARD = 1
    REVERSE = 2


class GenericRawMotorIndexes(IntEnum):
    LEFT_DRIVE = 0
    RIGHT_DRIVE = 1
    HEAD = 2
    LEG = 3


class GenericRawMotorModes(IntEnum):
    MOTOR_OFF = 0
    FORWARD = 1
    REVERSE = 2


class LinearVelocitySlewMethods(IntEnum):
    CONSTANT = 0
    PROPORTIONAL = 1


class MotorIndexes(IntEnum):
    LEFT_MOTOR_INDEX = 0
    RIGHT_MOTOR_INDEX = 1


class Drive(Commands):
    _did = 22

    @staticmethod
    def set_raw_motors(toy, left_mode: RawMotorModes, left_speed, right_mode: RawMotorModes, right_speed, proc=None):
        toy._execute(Drive._encode(toy, 1, proc, [left_mode, left_speed, right_mode, right_speed]))

    @staticmethod
    def reset_yaw(toy, proc=None):
        toy._execute(Drive._encode(toy, 6, proc))

    @staticmethod
    def drive_with_heading(toy, speed, heading, drive_flags: DriveFlags, proc=None):
        toy._execute(Drive._encode(toy, 7, proc, [speed, *to_bytes(heading, 2), drive_flags]))

    @staticmethod
    def generic_raw_motor(toy, index: GenericRawMotorIndexes, mode: GenericRawMotorModes, speed, proc=None):
        toy._execute(Drive._encode(toy, 11, proc, [index, mode, *speed]))

    @staticmethod
    def set_stabilization(toy, stabilization_index: StabilizationIndexes, proc=None):
        toy._execute(Drive._encode(toy, 12, proc, [stabilization_index]))

    @staticmethod
    def set_control_system_type(toy, s, s2, proc=None):  # unknown name
        toy._execute(Drive._encode(toy, 14, proc, [s, s2]))

    @staticmethod
    def set_pitch_torque_modification_value(toy, f, proc=None):  # Untested / Unknown Param Name
        toy._execute(Drive._encode(toy, 15, proc, [f]))

    @staticmethod
    def set_component_parameters(toy, s, s2, f_arr, proc=None):  # unknown name
        toy._execute(Drive._encode(toy, 32, proc, [s, s2, *struct.pack('>%df' % len(f_arr), *f_arr)]))

    @staticmethod
    def get_component_parameters(toy, s, s2, proc=None):  # unknown name
        data = toy._execute(Drive._encode(toy, 33, proc, [s, s2])).data
        return struct.unpack('>%df' % (len(data) // 4), data)

    @staticmethod
    def set_custom_control_system_timeout(toy, timeout, proc=None):
        toy._execute(Drive._encode(toy, 34, proc, to_bytes(timeout, 2)))

    @staticmethod
    def enable_motor_stall_notify(toy, enable, proc=None):
        toy._execute(Drive._encode(toy, 37, proc, [int(enable)]))

    motor_stall_notify = (22, 38, 0xff), lambda listener, p: listener(MotorStall(*p.data))

    @staticmethod
    def enable_motor_fault_notify(toy, enable, proc=None):
        toy._execute(Drive._encode(toy, 39, proc, [int(enable)]))

    motor_fault_notify = (22, 40, 0xff), lambda listener, p: listener(bool(p.data[0]))

    @staticmethod
    def get_motor_fault_state(toy, proc=None):
        return bool(toy._execute(Drive._encode(toy, 41, proc)).data[0])