import functools
import inspect
import logging
from typing import (
Any,
Callable,
ClassVar,
Dict,
NoReturn,
Optional,
Sequence,
Union,
overload,
)
import wpilib
from .magic_tunable import tunable
if wpilib.RobotBase.isSimulation():
getTime = wpilib.Timer.getFPGATimestamp
else:
from time import monotonic as getTime
[docs]
class IllegalCallError(TypeError):
pass
[docs]
class NoFirstStateError(ValueError):
pass
[docs]
class MultipleFirstStatesError(ValueError):
pass
[docs]
class MultipleDefaultStatesError(ValueError):
pass
[docs]
class InvalidStateName(ValueError):
pass
class _State:
def __init__(
self,
f: "StateMethod",
first: bool = False,
must_finish: bool = False,
*,
duration: Optional[float] = None,
is_default: bool = False,
) -> None:
name = f.__name__
# Can't define states that are named the same as things in the
# base class, will cause issues. Catch it early.
if hasattr(StateMachine, name):
raise InvalidStateName(f"cannot have a state named '{name}'")
# inspect the args, provide a correct call implementation
allowed_args = "self", "tm", "state_tm", "initial_call"
sig = inspect.signature(f)
args = []
invalid_args = []
for i, arg in enumerate(sig.parameters.values()):
if i == 0 and arg.name != "self":
raise ValueError(f"First argument to {name} must be 'self'")
if arg.kind is arg.VAR_POSITIONAL:
raise ValueError(f"Cannot use *args in signature for function {name}")
if arg.kind is arg.VAR_KEYWORD:
raise ValueError(
f"Cannot use **kwargs in signature for function {name}"
)
if arg.kind is arg.KEYWORD_ONLY:
raise ValueError(
f"Cannot use keyword-only parameters for function {name}"
)
if arg.name in allowed_args:
args.append(arg.name)
else:
invalid_args.append(arg.name)
if invalid_args:
raise ValueError(
"Invalid parameter names in %s: %s" % (name, ",".join(invalid_args))
)
self.name = name
self.description = inspect.getdoc(f)
self.first = first
self.must_finish = must_finish
self.is_default = is_default
self.duration = duration
varlist = {"f": f}
args_code = ",".join(args)
wrapper_code = f"lambda self, tm, state_tm, initial_call: f({args_code})"
self.run = eval(wrapper_code, varlist, varlist)
self.next_state: Optional[StateRef]
def __call__(self, *args, **kwargs) -> NoReturn:
raise IllegalCallError(
"Do not call states directly, use begin/next_state instead"
)
def __set_name__(self, owner: type, name: str) -> None:
# Don't allow aliasing of states, it probably won't do what users expect
if name != self.name:
raise InvalidStateName(
f"magicbot state '{self.name}' defined as attribute '{name}'"
)
if not issubclass(owner, StateMachine):
raise TypeError(f"magicbot state {name} defined in non-StateMachine")
# make durations tunable
if self.duration is not None:
duration_attr = name + "_duration"
# don't create it twice (in case of inheritance overriding)
if getattr(owner, duration_attr, None) is None:
setattr(
owner,
duration_attr,
tunable(self.duration, writeDefault=False, subtable="state"),
)
StateRef = Union[str, _State]
StateMethod = Callable[..., None]
class _StateData:
def __init__(self, wrapper: _State) -> None:
self.name = wrapper.name
self.duration_attr = "%s_duration" % self.name
self.expires: float = 0xFFFFFFFF
self.ran = False
self.run = wrapper.run
self.must_finish = wrapper.must_finish
if hasattr(wrapper, "next_state"):
self.next_state = wrapper.next_state
self.start_time: float
[docs]
def timed_state(
*,
duration: float,
next_state: Optional[StateRef] = None,
first: bool = False,
must_finish: bool = False,
) -> Callable[[StateMethod], _State]:
"""
If this decorator is applied to a function in an object that inherits
from :class:`.StateMachine`, it indicates that the function
is a state that will run for a set amount of time unless interrupted.
It is guaranteed that a timed_state will execute at least once, even if
it expires prior to being executed.
The decorated function can have the following arguments in any order:
- ``tm`` - The number of seconds since the state machine has started
- ``state_tm`` - The number of seconds since this state has been active
(note: it may not start at zero!)
- ``initial_call`` - Set to True when the state is initially called,
False otherwise. If the state is switched to multiple times, this
will be set to True at the start of each state execution.
:param duration: The length of time to run the state before progressing
to the next state
:param next_state: The name of the next state. If not specified, then
this will be the last state executed if time expires
:param first: If True, this state will be ran first
:param must_finish: If True, then this state will continue executing
even if ``engage()`` is not called. However,
if ``done()`` is called, execution will stop
regardless of whether this is set.
"""
def decorator(f: StateMethod) -> _State:
wrapper = _State(f, first, must_finish, duration=duration)
wrapper.next_state = next_state
return wrapper
return decorator
@overload
def state(
*,
first: bool = ...,
must_finish: bool = ...,
) -> Callable[[StateMethod], _State]: ...
@overload
def state(f: StateMethod) -> _State: ...
[docs]
def state(
f: Optional[StateMethod] = None,
*,
first: bool = False,
must_finish: bool = False,
) -> Union[Callable[[StateMethod], _State], _State]:
"""
If this decorator is applied to a function in an object that inherits
from :class:`.StateMachine`, it indicates that the function
is a state. The state will continue to be executed until the
``next_state`` function is executed.
The decorated function can have the following arguments in any order:
- ``tm`` - The number of seconds since the state machine has started
- ``state_tm`` - The number of seconds since this state has been active
(note: it may not start at zero!)
- ``initial_call`` - Set to True when the state is initially called,
False otherwise. If the state is switched to multiple times, this
will be set to True at the start of each state execution.
:param first: If True, this state will be ran first
:param must_finish: If True, then this state will continue executing
even if ``engage()`` is not called. However,
if ``done()`` is called, execution will stop
regardless of whether this is set.
"""
if f is None:
return lambda f: _State(f, first, must_finish)
return _State(f, first, must_finish)
[docs]
def default_state(f: StateMethod) -> _State:
"""
If this decorator is applied to a method in an object that inherits
from :class:`.StateMachine`, it indicates that the method
is a default state; that is, if no other states are executing, this
state will execute. If the state machine is always executing, the
default state will never execute.
There can only be a single default state in a StateMachine object.
The decorated function can have the following arguments in any order:
- ``tm`` - The number of seconds since the state machine has started
- ``state_tm`` - The number of seconds since this state has been active
(note: it may not start at zero!)
- ``initial_call`` - Set to True when the state is initially called,
False otherwise. If the state is switched to multiple times, this
will be set to True at the start of each state execution.
"""
return _State(f, first=False, must_finish=True, is_default=True)
def _get_class_members(cls: type) -> Dict[str, Any]:
"""Get the members of the given class in definition order, bases first."""
d = {}
for cls in reversed(cls.__mro__):
d.update(cls.__dict__)
return d
[docs]
class StateMachine:
'''
The StateMachine class is used to implement magicbot components that
allow one to easily define a `finite state machine (FSM)
<https://en.wikipedia.org/wiki/Finite-state_machine>`_ that can be
executed via the magicbot framework.
You create a component class that inherits from ``StateMachine``.
Each state is represented as a single function, and you indicate that
a function is a particular state by decorating it with one of the
following decorators:
* :func:`@default_state <.default_state>`
* :func:`@state <.state>`
* :func:`@timed_state <.timed_state>`
As the state machine executes, the decorated function representing the
current state will be called. Decorated state functions can receive the
following parameters (all of which are optional):
- ``tm`` - The number of seconds since autonomous has started
- ``state_tm`` - The number of seconds since this state has been active
(note: it may not start at zero!)
- ``initial_call`` - Set to True when the state is initially called,
False otherwise. If the state is switched to multiple times, this
will be set to True at the start of each state.
To be consistent with the magicbot philosophy, in order for the
state machine to execute its states you must call the :func:`engage`
function upon each execution of the main robot control loop. If you do
not call this function, then execution of the FSM will cease.
.. note:: If you wish for the FSM to continue executing state functions
regardless whether ``engage()`` is called, you must set the
``must_finish`` parameter in your state decorator to be True.
When execution ceases (because ``engage()`` was not called), the
:func:`done` function will be called and the FSM will be reset to the
starting state. The state functions will not be called again unless
``engage`` is called.
As a magicbot component, StateMachine contains an ``execute`` function that
will be called on each control loop. All state execution occurs from
within that function call. If you call other components from a
StateMachine, you should ensure that your StateMachine is declared
*before* the other components in your Robot class.
.. warning:: As StateMachine already contains an execute function,
there is no need to define your own ``execute`` function for
a state machine component -- if you override ``execute``,
then the state machine may not work correctly. Instead,
use the :func:`@default_state <.default_state>` decorator.
Here's a very simple example of how you might implement a shooter
automation component that moves a ball into a shooter when the
shooter is ready::
class ShooterAutomation(magicbot.StateMachine):
# Some other component
shooter: Shooter
ball_pusher: BallPusher
def fire(self):
"""This is called from the main loop."""
self.engage()
@state(first=True)
def begin_firing(self):
"""
This function will only be called IFF fire is called and
the FSM isn't currently in the 'firing' state. If fire
was not called, this function will not execute.
"""
self.shooter.enable()
if self.shooter.ready():
self.next_state('firing')
@timed_state(duration=1.0, must_finish=True)
def firing(self):
"""
Because must_finish=True, once the FSM has reached this state,
this state will continue executing even if engage isn't called.
"""
self.shooter.enable()
self.ball_pusher.push()
#
# Note that there is no execute function defined as part of
# this component
#
...
class MyRobot(magicbot.MagicRobot):
shooter_automation: ShooterAutomation
shooter: Shooter
ball_pusher: BallPusher
def teleopPeriodic(self):
if self.joystick.getTrigger():
self.shooter_automation.fire()
This object has a lot of really useful NetworkTables integration as well:
- tunables are created in /components/NAME/state
- state durations can be tuned here
- The 'current state' is output as it happens
- Descriptions and names of the states are here (for dashboard use)
.. warning:: This object is not intended to be threadsafe and should not
be accessed from multiple threads
'''
VERBOSE_LOGGING = False
#: A Python logging object automatically injected by magicbot.
#: It can be used to send messages to the log, instead of using print statements.
logger: logging.Logger
#: NT variable that indicates which state will be executed next (though,
#: does not guarantee that it will be executed). Will return an empty
#: string if the state machine is not currently engaged.
current_state = tunable("", subtable="state")
state_names: ClassVar[tunable[Sequence[str]]]
state_descriptions: ClassVar[tunable[Sequence[str]]]
def __new__(cls) -> "StateMachine":
# choose to use __new__ instead of __init__
o = super().__new__(cls)
o._build_states()
return o
# TODO: when this gets invoked, tunables need to be setup on
# the object first
def _build_states(self) -> None:
has_first = False
# problem: the user interface won't know which entries are the
# current variables being used by the robot. So, we setup
# an array with the names, and the dashboard uses that
# to determine the ordering too
nt_names = []
nt_desc = []
states = {}
cls = type(self)
default_state = None
# for each state function:
for name, state in _get_class_members(cls).items():
if not isinstance(state, _State):
continue
# is this the first state to execute?
if state.first:
if has_first:
raise MultipleFirstStatesError(
"Multiple states were specified as the first state!"
)
self.__first = name
has_first = True
state_data = _StateData(state)
states[name] = state_data
nt_names.append(name)
nt_desc.append(state.description or "")
if state.is_default:
if default_state is not None:
raise MultipleDefaultStatesError(
"Multiple default states are not allowed"
)
default_state = state_data
if not has_first:
raise NoFirstStateError(
"Starting state not defined! Use first=True on a state decorator"
)
# NOTE: this depends on tunables being bound after this function is called
cls.state_names = tunable(nt_names, subtable="state")
cls.state_descriptions = tunable(nt_desc, subtable="state")
# Indicates that an external party wishes the state machine to execute
self.__should_engage = False
# Indicates that the state machine is currently executing
self.__engaged = False
# A dictionary of states
self.__states = states
# The currently executing state, or None if not executing
self.__state: Optional[_StateData] = None
# The default state
self.__default_state = default_state
# Variable to store time in
self.__start = 0
@property
def is_executing(self) -> bool:
""":returns: True if the state machine is executing states"""
# return self.__state is not None
return self.__engaged
[docs]
def on_enable(self) -> None:
"""
magicbot component API: called when autonomous/teleop is enabled
"""
pass
[docs]
def on_disable(self) -> None:
"""
magicbot component API: called when autonomous/teleop is disabled
"""
self.done()
[docs]
def engage(
self,
initial_state: Optional[StateRef] = None,
force: bool = False,
) -> None:
"""
This signals that you want the state machine to execute its
states.
:param initial_state: If specified and execution is not currently
occurring, start in this state instead of
in the 'first' state
:param force: If True, will transition even if the state
machine is currently active.
"""
self.__should_engage = True
if force or self.__state is None or self.__state is self.__default_state:
if initial_state:
self.next_state(initial_state)
else:
self.next_state(self.__first)
[docs]
def next_state(self, state: StateRef) -> None:
"""Call this function to transition to the next state
:param state: Name of the state to transition to
.. note:: This should only be called from one of the state functions
"""
if isinstance(state, _State):
state = state.name
state_data = self.__states[state]
state_data.ran = False
self.current_state = state
self.__state = state_data
[docs]
def next_state_now(self, state: StateRef) -> None:
"""Call this function to transition to the next state, and call the next
state function immediately. Prefer to use :meth:`next_state` instead.
:param state: Name of the state to transition to
.. note:: This should only be called from one of the state functions
"""
self.next_state(state)
# TODO: may want to do this differently?
self.execute()
[docs]
def done(self) -> None:
"""Call this function to end execution of the state machine.
This function will always be called when a state machine ends. Even if
the engage function is called repeatedly, done() will be called.
.. note:: If you wish to do something each time execution ceases,
override this function (but be sure to call
``super().done()``!)
"""
if self.VERBOSE_LOGGING and self.__state is not None:
self.logger.info("Stopped state machine execution")
self.__state = None
self.__engaged = False
self.current_state = ""
[docs]
def execute(self) -> None:
"""
magicbot component API: This is called on each iteration of the
control loop. Most of the time, you will not want to override
this function. If you find you want to, you may want to use the
@default_state mechanism instead.
"""
now = getTime()
if not self.__engaged:
if self.__should_engage:
self.__start = now
self.__engaged = True
elif self.__default_state is None:
return
# tm is the number of seconds that the state machine has been executing
tm = now - self.__start
state = self.__state
done_called = False
# we adjust this so that if we have states chained together,
# then the total time it runs is the amount of time of the
# states. Otherwise, the time drifts.
new_state_start = tm
# determine if the time has passed to execute the next state
# -> intentionally comes first
if state is not None and state.ran and state.expires < tm:
new_state_start = state.expires
if state.next_state is None:
# If the state expires and it's the last state, if the machine
# is still engaged then it should cycle back to the beginning
# ... but we should call done() first
done_called = True
self.done()
if self.__should_engage:
self.next_state(self.__first)
state = self.__state
else:
state = None
else:
self.next_state(state.next_state)
state = self.__state
# deactivate the current state unless engage was called or
# must_finish was set
if not (self.__should_engage or state is not None and state.must_finish):
state = None
# if there is no state to execute and there is a default
# state, do the default state
if state is None and self.__default_state is not None:
state = self.__default_state
if self.__state != state:
state.ran = False
self.__state = state
if state is not None:
# is this the first time this was executed?
initial_call = not state.ran
if initial_call:
state.ran = True
state.start_time = new_state_start
state.expires = new_state_start + getattr(
self, state.duration_attr, 0xFFFFFFFF
)
if self.VERBOSE_LOGGING:
self.logger.info("%.3fs: Entering state: %s", tm, state.name)
# execute the state function, passing it the arguments
state.run(self, tm, tm - state.start_time, initial_call)
elif not done_called:
# or clear the state
self.done()
# Reset this each time
self.__should_engage = False
[docs]
class AutonomousStateMachine(StateMachine):
"""
This is a specialized version of the StateMachine that is designed
to be used as an autonomous mode. There are a few key differences:
- The :func:`.engage` function is always called, so the state machine
will always run to completion unless done() is called
- VERBOSE_LOGGING is set to True, so a log message will be printed out upon
each state transition
"""
VERBOSE_LOGGING = True
[docs]
def on_enable(self) -> None:
super().on_enable()
self.__engaged = True
[docs]
def on_iteration(self, tm: float) -> None:
# TODO, remove the on_iteration function in 2017?
# Only engage the state machine until its execution finishes, otherwise
# it will just keep repeating
#
# This is because if you keep calling engage(), the state machine will
# loop. I'm tempted to change that, but I think it would lead to unexpected
# side effects. Will have to contemplate this...
if self.__engaged:
self.engage()
self.execute()
self.__engaged = self.is_executing
[docs]
def done(self) -> None:
super().done()
self._StateMachine__should_engage = False
self.__engaged = False