Raw Text Content QR
fish-zebra-ant



# Interactively identify printer axis motors via web ui prompts
#
# Copyright (C) 2026  the Klipper contributors
#
# This file may be distributed under the terms of the GNU GPLv3 license.
import re

AXES = "xyz"

STEPPER_NAME_R = re.compile(r'(?i)^stepper_([a-z])(\d*)$')

# Stepper pins that identify which motor/driver socket drives the rail.
# When the changes are written through the SAVE_CONFIG mechanism these
# pins are swapped between the stepper sections (the sections keep their
# names, so all changes end up in the "#*#" block at the end of the
# config file)
MOTOR_OPTIONS = ('step_pin', 'dir_pin', 'enable_pin')


class AxisDiscovery:
    def __init__(self, config):
        self.printer = printer = config.get_printer()
        self.gcode = printer.lookup_object('gcode')
        self.wizard = None
        # Register g-code commands
        handlers = ['AXIS_DISCOVERY', 'AXIS_DISCOVERY_NEXT',
                    'AXIS_DISCOVERY_REPEAT', 'AXIS_DISCOVERY_ANSWER',
                    'AXIS_DISCOVERY_APPLY', 'AXIS_DISCOVERY_CANCEL']
        for cmd in handlers:
            func = getattr(self, 'cmd_' + cmd)
            desc = getattr(self, 'cmd_' + cmd + '_help', None)
            self.gcode.register_command(cmd, func, desc=desc)
        printer.register_event_handler("klippy:disconnect",
                                       self._handle_disconnect)

    def _handle_disconnect(self):
        self.wizard = None

    # Prompt emission (compatible with the web ui "action:prompt_*" protocol)
    def _prompt_action(self, action, msg=None):
        self.gcode.respond_raw("// action:prompt_%s%s"
                               % (action, "" if msg is None else " " + msg))

    def _prompt_begin(self, headline):
        self._prompt_action("begin", headline)

    def _prompt_text(self, msg):
        self._prompt_action("text", msg)

    def _prompt_button(self, msg):
        self._prompt_action("button", msg)

    def _prompt_footer_button(self, msg):
        self._prompt_action("footer_button", msg)

    def _prompt_show(self):
        self._prompt_action("show")

    def _prompt_end(self):
        self._prompt_action("end")

    # Wizard helpers
    def _check_not_printing(self):
        vsd = self.printer.lookup_object('virtual_sdcard', None)
        if vsd is not None and vsd.is_active():
            raise self.printer.command_error(
                "Command not available during a print")

    def _get_wizard(self, gcmd):
        if self.wizard is None:
            raise gcmd.error(
                "No axis discovery in progress (run AXIS_DISCOVERY first)")
        return self.wizard

    def _lookup_motors(self):
        toolhead = self.printer.lookup_object('toolhead')
        kin = toolhead.get_kinematics()
        if kin.__class__.__name__ != 'CartKinematics':
            raise self.printer.command_error(
                "AXIS_DISCOVERY supports [cartesian] kinematics only")
        rails = getattr(kin, 'rails', None)
        if rails is None or len(rails) != 3:
            raise self.printer.command_error(
                "AXIS_DISCOVERY is not supported with a dual carriage")
        raw_config = self.printer.lookup_object('configfile').get_status(
            0.)['config']
        motors = []
        for rail, axis in zip(rails, AXES):
            for mcu_stepper in rail.get_steppers():
                name = mcu_stepper.get_name()
                m = STEPPER_NAME_R.match(name)
                if m is None or m.group(1).lower() != axis:
                    raise self.printer.command_error(
                        "Unexpected stepper section '%s'" % (name,))
                dir_pin = raw_config.get(name, {}).get('dir_pin', '')
                motors.append({
                    'name': name, 'axis': axis, 'suffix': m.group(2),
                    'inverted': str(dir_pin).strip().startswith('!'),
                    'answer': None,
                })
        return motors

    def _move_current(self, wizard):
        self._check_not_printing()
        motor = wizard['motors'][wizard['index']]
        force_move = self.printer.lookup_object('force_move')
        stepper_enable = self.printer.lookup_object('stepper_enable')
        mcu_stepper = force_move.lookup_stepper(motor['name'])
        did_enable = stepper_enable.set_motors_enable([motor['name']], True)
        try:
            force_move.manual_move(mcu_stepper, wizard['distance'],
                                   wizard['speed'], wizard['accel'])
        finally:
            if did_enable:
                stepper_enable.set_motors_enable([motor['name']], False)

    def _show_question(self, wizard):
        motor = wizard['motors'][wizard['index']]
        self._prompt_begin("Axis Discovery (%d/%d)"
                           % (wizard['index'] + 1, len(wizard['motors'])))
        self._prompt_text("Moving %s in the positive direction..."
                          % (motor['name'],))
        self._prompt_text("Which motor position moved and in which"
                          " direction?")
        for axis in AXES:
            slots = [m for m in wizard['motors'] if m['axis'] == axis]
            slots.sort(key=lambda m: (len(m['suffix']), m['suffix']))
            for slot in slots:
                short = axis.upper() + slot['suffix']
                slot_param = ""
                if slot['suffix']:
                    slot_param = " SLOT=%s" % (slot['suffix'],)
                self._prompt_action("button_group_start")
                self._prompt_button(
                    "%s+|AXIS_DISCOVERY_ANSWER I=%d AXIS=%s%s DIR=+"
                    % (short, wizard['index'], axis, slot_param))
                self._prompt_button(
                    "%s-|AXIS_DISCOVERY_ANSWER I=%d AXIS=%s%s DIR=-"
                    % (short, wizard['index'], axis, slot_param))
                self._prompt_action("button_group_end")
        self._prompt_footer_button("Move again|AXIS_DISCOVERY_REPEAT")
        self._prompt_footer_button("Cancel|AXIS_DISCOVERY_CANCEL")
        self._prompt_show()

    def _calc_changes(self, wizard):
        motors = wizard['motors']
        # Determine the target motor position of each motor.  Answers
        # without an explicit SLOT are assigned to the remaining free
        # positions of their axis.
        taken = set()
        targets = [None] * len(motors)
        for i, motor in enumerate(motors):
            axis, suffix, sign = motor['answer']
            if suffix is not None:
                targets[i] = (axis, suffix)
                taken.add(targets[i])
        for i, motor in enumerate(motors):
            if targets[i] is not None:
                continue
            axis = motor['answer'][0]
            free = sorted(
                [m['suffix'] for m in motors
                 if m['axis'] == axis and (axis, m['suffix']) not in taken],
                key=lambda s: (len(s), s))
            if not free:
                raise self.printer.command_error(
                    "Axis %s has no free motor positions left"
                    % (axis.upper(),))
            targets[i] = (axis, free[0])
            taken.add(targets[i])
        rename_map = {}
        invert_new = set()
        info = []
        for i, motor in enumerate(motors):
            t_axis, t_suffix = targets[i]
            sign = motor['answer'][2]
            new_name = "stepper_%s%s" % (t_axis, t_suffix)
            label = motor['name']
            if new_name != motor['name']:
                rename_map[motor['name'].lower()] = new_name.lower()
                label += " -> " + new_name
            if sign < 0:
                invert_new.add(new_name.lower())
                label += " (dir inverted)"
            info.append(label)
        return rename_map, invert_new, info

    def _show_summary(self, wizard):
        rename_map, invert_new, info = self._calc_changes(wizard)
        self._prompt_begin("Axis Discovery - confirm changes")
        if not rename_map and not invert_new:
            self._prompt_text("No changes required.")
        else:
            for line in info:
                self._prompt_text(line)
            self._prompt_text("The changes are saved to the SAVE_CONFIG"
                              " block at the end of the config file.")
            self._prompt_footer_button(
                "Apply & Restart|AXIS_DISCOVERY_APPLY")
        self._prompt_footer_button("Cancel|AXIS_DISCOVERY_CANCEL")
        self._prompt_show()

    # G-code commands
    cmd_AXIS_DISCOVERY_help = "Interactively identify axis motors"
    def cmd_AXIS_DISCOVERY(self, gcmd):
        self._check_not_printing()
        distance = gcmd.get_float('DISTANCE', 10., above=0., maxval=100.)
        speed = gcmd.get_float('SPEED', 25., above=0., maxval=200.)
        accel = gcmd.get_float('ACCEL', 0., minval=0.)
        motors = self._lookup_motors()
        # Manual stepper moves invalidate the toolhead position
        toolhead = self.printer.lookup_object('toolhead')
        toolhead.get_kinematics().clear_homing_state(AXES)
        self.wizard = {'motors': motors, 'index': -1, 'spec': {},
                       'unspecified': {},
                       'distance': distance, 'speed': speed, 'accel': accel}
        self._prompt_begin("Axis Discovery")
        self._prompt_text("Detected motors:")
        for motor in motors:
            self._prompt_text("- %s (axis %s, dir_pin %s)"
                              % (motor['name'], motor['axis'].upper(),
                                 "inverted" if motor['inverted']
                                 else "normal"))
        self._prompt_text("Each motor is moved %.1fmm in the positive"
                          " direction. Watch which carriage moves and where"
                          " it goes." % (distance,))
        self._prompt_text("Homing is invalidated, the printer must be"
                          " re-homed after discovery.")
        self._prompt_footer_button("Start|AXIS_DISCOVERY_NEXT I=0")
        self._prompt_footer_button("Cancel|AXIS_DISCOVERY_CANCEL")
        self._prompt_show()

    cmd_AXIS_DISCOVERY_NEXT_help = "Move the next motor and request input"
    def cmd_AXIS_DISCOVERY_NEXT(self, gcmd):
        wizard = self._get_wizard(gcmd)
        expected = gcmd.get_int('I', None)
        if expected is not None and expected != wizard['index'] + 1:
            # Stale prompt button (eg, clicked twice); ignore it
            return
        if wizard['index'] + 1 >= len(wizard['motors']):
            self._show_summary(wizard)
            return
        wizard['index'] += 1
        self._move_current(wizard)
        self._show_question(wizard)

    cmd_AXIS_DISCOVERY_REPEAT_help = "Repeat the move of the current motor"
    def cmd_AXIS_DISCOVERY_REPEAT(self, gcmd):
        wizard = self._get_wizard(gcmd)
        if wizard['index'] < 0 or wizard['index'] >= len(wizard['motors']):
            raise gcmd.error("No motor is currently being identified")
        self._move_current(wizard)
        self._show_question(wizard)

    cmd_AXIS_DISCOVERY_ANSWER_help = \
        "Report the observed motor position and direction"
    def cmd_AXIS_DISCOVERY_ANSWER(self, gcmd):
        wizard = self._get_wizard(gcmd)
        expected = gcmd.get_int('I', None)
        if expected is not None and expected != wizard['index']:
            raise gcmd.error(
                "Stale axis discovery answer (motor index %d expected)"
                % (wizard['index'],))
        if wizard['index'] < 0 or wizard['index'] >= len(wizard['motors']):
            raise gcmd.error("No motor is awaiting an answer")
        motor = wizard['motors'][wizard['index']]
        axis = gcmd.get('AXIS').lower()
        if axis not in AXES:
            raise gcmd.error("Invalid AXIS '%s'" % (axis,))
        suffix = gcmd.get('SLOT', None)
        if suffix is not None and (axis, suffix) not in [
                (m['axis'], m['suffix']) for m in wizard['motors']]:
            raise gcmd.error("Invalid SLOT '%s' for axis %s"
                             % (suffix, axis.upper()))
        direction = gcmd.get('DIR')
        if direction == '+':
            sign = 1
        elif direction == '-':
            sign = -1
        else:
            raise gcmd.error("Invalid DIR '%s' (must be + or -)"
                             % (direction,))
        # Check that the target motor position is still free and that the
        # axis does not receive more motors than it has positions
        used = len([1 for (a, s) in wizard['spec'] if a == axis]) \
            + wizard['unspecified'].get(axis, 0)
        avail = len([1 for m in wizard['motors'] if m['axis'] == axis])
        if used >= avail:
            raise gcmd.error(
                "Axis %s has no free motor positions left"
                % (axis.upper(),))
        if suffix is None:
            wizard['unspecified'][axis] = \
                wizard['unspecified'].get(axis, 0) + 1
        else:
            if (axis, suffix) in wizard['spec']:
                raise gcmd.error(
                    "Motor stepper_%s%s has already been assigned"
                    % (axis, suffix))
            wizard['spec'][(axis, suffix)] = motor
        motor['answer'] = (axis, suffix, sign)
        wizard['index'] += 1
        if wizard['index'] < len(wizard['motors']):
            self._move_current(wizard)
            self._show_question(wizard)
        else:
            self._show_summary(wizard)

    cmd_AXIS_DISCOVERY_APPLY_help = "Apply axis discovery changes and restart"
    def cmd_AXIS_DISCOVERY_APPLY(self, gcmd):
        wizard = self._get_wizard(gcmd)
        for motor in wizard['motors']:
            if motor['answer'] is None:
                raise gcmd.error("Motor %s has not been identified yet"
                                 % (motor['name'],))
        self._check_not_printing()
        rename_map, invert_new, info = self._calc_changes(wizard)
        if not rename_map and not invert_new:
            self.wizard = None
            self._prompt_end()
            gcmd.respond_info("Axis discovery: no changes required")
            return
        self._prompt_end()
        # The stepper sections keep their names and the motor/driver pins
        # are swapped between them, so all changes end up in the "#*#"
        # block at the end of the config file
        touched = self._apply_autosave(rename_map, invert_new,
                                       wizard['motors'])
        self.wizard = None
        gcmd.respond_info("Axis discovery: changes saved to the"
                          " SAVE_CONFIG block, restarting...")
        try:
            self.gcode.run_script_from_command("SAVE_CONFIG")
        except self.printer.command_error as e:
            # Do not leave pending entries behind (they would otherwise be
            # applied by any later SAVE_CONFIG command)
            self._discard_autosave(touched)
            raise gcmd.error("Unable to save via SAVE_CONFIG: %s" % (e,))

    cmd_AXIS_DISCOVERY_CANCEL_help = "Abort the axis discovery wizard"
    def cmd_AXIS_DISCOVERY_CANCEL(self, gcmd):
        self.wizard = None
        self._prompt_end()
        gcmd.respond_info("Axis discovery cancelled (re-home the printer)")

    # SAVE_CONFIG support
    def _source_of(self, rename_map, motors):
        # Map each stepper section to the section whose motor/driver
        # options it receives (sections keep their names)
        source_of = {}
        for motor in motors:
            new_name = rename_map.get(motor['name'].lower(), motor['name'])
            source_of[new_name.lower()] = motor['name']
        return source_of

    def _raw_section(self, raw, name):
        if name in raw:
            return raw[name]
        return raw.get(name.lower(), {})

    def _tmc_sections(self, raw, sec):
        out = {}
        for key, opts in raw.items():
            parts = key.split(' ', 1)
            if (len(parts) == 2 and parts[0].lower().startswith('tmc')
                    and parts[1].strip().lower() == sec.lower()):
                out[parts[0]] = opts
        return out

    def _autosave_ready(self, rename_map, invert_new, motors):
        # Returns None when the changes can be written through the
        # SAVE_CONFIG mechanism, otherwise the reason they cannot
        cfgname = self.printer.get_start_args()['config_file']
        files = {}
        try:
            self._read_config_file(cfgname, files, set())
        except self.printer.command_error as e:
            return str(e)
        main_path = os.path.abspath(cfgname)
        include_sections = set()
        for fname, lines in files.items():
            if fname == main_path:
                continue
            for line in lines:
                hm = SECTION_HEADER_R.match(line)
                if hm is not None:
                    include_sections.add(hm.group(2).strip().lower())
        raw = self.printer.lookup_object('configfile').get_status(
            0.)['config']
        for slot, src in self._source_of(rename_map, motors).items():
            for name in set([slot, src]):
                if name.lower() in include_sections:
                    return ("section [%s] is defined in an included config"
                            " file" % (name,))
            ssec = self._raw_section(raw, slot)
            osec = self._raw_section(raw, src)
            for opt in MOTOR_OPTIONS:
                if (opt in ssec) != (opt in osec):
                    return ("option '%s' is not defined in both [%s] and"
                            " [%s]" % (opt, slot, src))
            s_tmc = self._tmc_sections(raw, slot)
            o_tmc = self._tmc_sections(raw, src)
            if sorted(s_tmc) != sorted(o_tmc):
                return ("[%s] and [%s] do not use the same tmc driver"
                        " sections" % (slot, src))
            for chip in s_tmc:
                if sorted(s_tmc[chip]) != sorted(o_tmc[chip]):
                    return ("[%s %s] and [%s %s] do not define the same"
                            " options" % (chip, slot, chip, src))
        return None

    def _apply_autosave(self, rename_map, invert_new, motors):
        pconfig = self.printer.lookup_object('configfile')
        raw = pconfig.get_status(0.)['config']
        touched = []
        for slot, src in self._source_of(rename_map, motors).items():
            ssec = self._raw_section(raw, slot)
            osec = self._raw_section(raw, src)
            invert = slot in invert_new
            for opt in MOTOR_OPTIONS:
                if opt not in osec:
                    continue
                val = str(osec[opt]).strip()
                if opt == 'dir_pin' and invert:
                    val = self._invert_pin(val, slot)
                if val == str(ssec.get(opt, '')).strip():
                    # No change necessary; keep the config file untouched
                    continue
                pconfig.set(slot, opt, val)
                touched.append((slot, opt))
            s_tmc = self._tmc_sections(raw, slot)
            o_tmc = self._tmc_sections(raw, src)
            for chip, opts in o_tmc.items():
                for opt, val in opts.items():
                    if str(val).strip() == str(
                            s_tmc.get(chip, {}).get(opt, '')).strip():
                        continue
                    pconfig.set("%s %s" % (chip, slot), opt, str(val).strip())
                    touched.append(("%s %s" % (chip, slot), opt))
        return touched

    def _discard_autosave(self, touched):
        # Remove the entries set by a SAVE_CONFIG command that failed, so
        # that they do not affect any later SAVE_CONFIG command
        pconfig = self.printer.lookup_object('configfile')
        fileconfig = pconfig.fileconfig
        sections = set([section for section, option in touched])
        for section, option in touched:
            if fileconfig.has_section(section):
                fileconfig.remove_option(section, option)
                if not fileconfig.options(section):
                    fileconfig.remove_section(section)
        pending = dict(pconfig.status_save_pending)
        for section in sections:
            if fileconfig.has_section(section):
                pending[section] = {
                    opt: fileconfig.get(section, opt)
                    for opt in fileconfig.options(section)}
            else:
                pending.pop(section, None)
        pconfig.status_save_pending = pending
        pconfig.save_config_pending = bool(pending)

    def _invert_pin(self, pin, section):
        pin = pin.strip()
        if not pin:
            raise self.printer.command_error(
                "Empty dir_pin in section '%s'" % (section,))
        if pin.startswith('!'):
            return pin[1:]
        return '!' + pin



def load_config(config):
    return AxisDiscovery(config)

Read 4 times, last 4 days ago